Skip to content

Contributing Guide

Everyone can contribute to Ocean Emulator, and we value everyone's contributions. There are several ways to help, including:

  • Reporting bugs or feature requests in our issue tracker.
  • Contributing PRs to our code base.
  • Writing or editing documentation. (Yes, typo fixes are welcome!)

This project follows the M2LInES Code of Conduct.

Contributing code with pull requests

TL;DR 
git clone git@github.com:Open-Athena/Ocean_Emulator.git
cd Ocean_Emulator
uv sync --dev
source .venv/bin/activate
uvx pre-commit install
uvx pre-commit run --all-files # also creates config schemas for validation (see below)

# dev
uv run pytest -m "not manual and not cuda"
uv run pytest --benchmark-only --benchmark-autosave
uv run pytest-benchmark compare 0001 0002

# push new remote branch to make a PR
git push -u origin <feature-branch>

# sync branch
git pull origin main --rebase
git push --force-with-lease

  1. (If you're not a core maintainer), please fork the repository by clicking the Fork button on the repository page.

  2. Clone the repository (via ssh recommended!) and change into the root directory. 

    # if you're using a fork, make sure to clone your fork's repo
git clone https://github.com/Open-Athena/Ocean_Emulator.git
# preferred method, but requires setting up an ssh key with Github.
git clone git@github.com:Open-Athena/Ocean_Emulator.git
# or, using the Github CLI
gh repo clone Open-Athena/Ocean_Emulator

# then, change directory
cd Ocean_Emulator

  3. Install developer dependencies using uv: 

    uv sync --dev
    Then, activate the environment that uv created: 
    source .venv/bin/activate

  4. (If forked) Add the original repository as an upstream remote, so you can sync your changes. 

    # via http
git remote add upstream https://github.com/Open-Athena/Ocean_Emulator.git
# via ssh
git remote add upstream git@github.com:Open-Athena/Ocean_Emulator.git

  5. Check out feature branches where you will develop from: 

    git checkout -b name-of-change

  6. Perform project lifecycle routines as needed during development: 

    # run tests
uvx pytest -m "not manual and not cuda"

For more details on how to run specific tests, please see the section below.

To validate the PhysicsNeMo-based container locally: 

scripts/container/build_physicsnemo_25_11.sh
scripts/container/run_cuda_tests_in_image.sh
The corresponding CI workflow is Container PhysicsNeMo 25.11 in .github/workflows/container-physicsnemo.yml (x86 build + smoke checks, publish, and containerized CPU/GPU tests).

Recommended: For convenience, we've collected lint checks as a pre-commit hook.

To install the pre-commit hook (which will run before every commit), call: 

uvx pre-commit install

To run all checks manually, you can run: 

# run against staged files
uvx pre-commit run
# run against all files in the project
uvx pre-commit run --all-files

If you want to commit without running pre-commit checks, you're always free to use the --no-verify flag: 

git commit --no-verfiy -m "WIP"

Sometimes, you may want to skip just one check, but run the rest of the pre-commit. You can accomplish this by setting an environment variable: 

export SKIP=ruff
uvx pre-commit run

It's totally ok to make lots of small commits as you develop your feature! Please, make sure to write commit messages along the way.

Sometimes, you may change code that changes the project's performance characteristics. To measure the code's current local performance, you can run:

uv run pytest --benchmark-only --benchmark-autosave

This will output stats on all our current benchmarks performance in your terminal. To compare this performance to previous changes, run: 

uv run pytest-benchmark compare <id1> <id2>
# For example, comparing change 1 to change 2:
uv run pytest-benchmark compare 0001 0002

Please see the Benchmarks & Profiling section below for more details.

  1. Before submitting a pull request, please sync with the main repo via rebase: 
    git pull origin main --rebase
# if working in a fork
git pull upstream main --rebase

If the rebase requires that you force push to your remote feature branch, we recommend using --force-with-lease: 

git push --force-with-lease

  1. Finally, when you're ready to submit a pull request — say, when all checks have passed — push your change on your development branch so you can create a pull request: 

    git push -u origin name-of-change

  2. Before you make the final merge, please make sure your commits are tidy and well-named. To do this, you can either use the Squash and merge button to commit (the default), or merge commits after you've cleaned up the commit history. For this, we recommend performing an interactive rebase: 

    git rebase -i <starting-commit>

  3. Celebrate submitting your patch to Ocean Emulator — well done!

Running Ocean Emulator

Please review the How to Run guide to learn how to train, evaluate, and visualize our emulators. To learn how to customize experiments and hyperparameters, please read our configuration guide.

VS Code Integration

If you're using VS Code, we recommend installing the ruff and mypy extensions. For the latter, you'll want to configure it to use pyproject.toml, which you can do with a .vscode/settings.json file:

{
    "mypy-type-checker.args": [
        "\"--config=pyproject.toml\""
    ]
}

Testing Ocean Emulator

TL;DR 
# local dev / CI
pytest -m "not manual and not cuda"
# with more CPU cores
pytest -m "not manual and not cuda" -n auto
# all manual tests
pytest -m manual

We use pytest as a test runner. All tests in this project have several marks that allow developers to control what tests are run locally. Two marks of particular interest are cuda and manual.

"cuda" tests are tests that require an NVIDIA GPU to run. If tests use the device fixture, then they'll automatically be configured to run on both GPU and CPU simultaneously. Certain tests, however, can be marked with @pytest.mark.cuda if they need to run on that hardware. To run CUDA-only tests, call:

pytest -m cuda
And to exclude all cuda-marked tests, run:

pytest -m "not cuda"

"manual" tests are not run in continuous integration (CI), but are useful checks during the development process. For example, evaluating if two model weights are equal is marked manual. All manual tests can be run like so:

pytest -m manual

To exclude manual tests, run:

pytest -m "not manual"

To run the same tests that are run in CI, please run the following:

pytest -m "not manual and not cuda"

Running tests in the PhysicsNeMo container

If you want to run tests using the PhysicsNeMo 25.11-based image, build it first:

BUILD_APPTAINER=0 scripts/container/build_physicsnemo_25_11.sh

Run the CPU test set (same marker expression as CPU CI) inside the built image:

docker run --rm \
  -v "$PWD":/repo \
  -w /workspace \
  ocean-emulator:physicsnemo-25.11 \
  bash -lc '. .venv/bin/activate && cd /repo && python -m pytest -m "not manual and not cuda"'

Run CUDA tests in the built image:

scripts/container/run_cuda_tests_in_image.sh

You can pass custom pytest arguments for CUDA runs, for example:

PYTEST_ARGS="-k test_trainer" scripts/container/run_cuda_tests_in_image.sh

Testing with Multitons

We have a set of singletons in the code which use the "Multiton" helper to prevent tests from interfering with each other. When writing tests, you can either:

   def test_foo():
      with MultitonScope():
         # set up whatever singletons you need
         Normalize.init_instance(...)
         assert ...

Or you can initialize them in a Generator-based fixture:

   @pytest.fixture()
   def my_fixture():
       with MultitonScope():
           Normalize.init_instance(...)
           yield

   def test_foo(my_fixture):
       assert ... # in this code, the Normalize instance is the one from my_fixture

Preventing checking-in secrets

In our pre-commit check, we use a tool developed by Yelp that detects strings that look suspiciously like secrets and raises alarms. If this is blocking your patch and you've manually inspected the sources for secrets and vetted that there are, in fact, none checked in, the following command will regenerate a metadata file to pass this check:

uvx detect-secrets scan > .secrets.baseline

Please check in the baseline after generating.

Benchmarking & Profiling

We use pytest-benchmark to measure performance regressions in this project. Our intentions are to cultivate a culture of writing performant programs. To this end, we offer users the following tools:

To run local benchmarks and save their status locally (associated with the current commit), execute:

uv run pytest --benchmark-only --benchmark-autosave

To compare benchmark run 0001 to 0002, you can run:

uv run pytest-benchmark compare 0001 0002

Please check your local .benchmarks/ directory to see other benchmarks runs for comparison.

To generate a histogram plot of several local benchmark runs, you may use the --histogram=FILENAME-PREFIX flag:

uv run pytest-benchmark compare 'Darwin-CPython-3.10-64bit/*' --histogram

Instead of merely benchmarking performance, sometimes you may want to inspect the details of how benchmarks run. This is useful, for example, for white-box performance optimization. To collect a cProfile trace for each benchmark, run:

uv run pytest --benchmark-only --benchmark-cprofile="tottime_per" --benchmark-cprofile-dump

(Please consult these docs to see all available values for the --benchmark-cprofile flag.)

This will generate a .prof file located in the . directory (by default). You can visualize this trace with snakeviz like so:

uvx snakeviz <benchmark_path>.prof

Profiling CPU Usage + Memory

We also have a few other profiling tools available in the environment, including:

py-spy, which captures python + native CPU usage: 

uv run py-spy record --native -o profile.svg -- ./.venv/bin/python  src/ocean_emulators/train.py configs/samudra_vnext/train.yaml

memray, which captures peak memory usage:

uv run memray run src/ocean_emulators/train.py --config configs/samudra_vnext/train.yaml
uv run memray flamegraph path/to/memray-output.bin

And scalene, which shows per-line python/native CPU usage, memory usage and GPU (though the latter is a bit deceptive since it is async wrt the highlighted code).

uv run scalene src/ocean_emulators/train.py configs/samudra_vnext/train.yaml

Profiling CUDA Memory

You can turn on profiling of CUDA memory by setting the profiler.cuda_snapshot_frequency to a non-None value in the config. eg:

uv run memray run src/ocean_emulators/train.py --config configs/samudra_vnext/train.yaml --profiler.cuda_snapshot_frequency 10

This will take a snapshot of the CUDA memory every 10 batches in the output directory. These can be visualized with https://docs.pytorch.org/memory_viz -- see https://pytorch.org/blog/understanding-gpu-memory-1/ for more details.

Data Engineering

To learn about how to access or recreate our datasets, please review our data guide.