Skip to content

Add sharded ocean climate simulation test#274

Open
glwagner wants to merge 3 commits intomainfrom
gw/sharded-ocean-climate-test
Open

Add sharded ocean climate simulation test#274
glwagner wants to merge 3 commits intomainfrom
gw/sharded-ocean-climate-test

Conversation

@glwagner
Copy link
Copy Markdown
Collaborator

Summary

  • Add compile and run scripts for sharded ocean climate simulation using a tripolar grid with the coupled ocean-atmosphere model (sharding/sharded_ocean_climate_simulation_{run,compile}.jl)
  • Refactor data_free_ocean_climate_model_init to accept (arch, Nx, Ny, Nz) with grid_type, halo, Δt, and set_initial_conditions kwargs
  • Add simple_tripolar_grid helper in model_utils.jl
  • Update Compile.yml and Run.yml workflows to include sharded_ocean_climate in the sharded test matrices

This tests sharding with a different grid topology (tripolar vs lat-lon) and atmosphere coupling, complementing the existing sharded_baroclinic_instability test which only uses a simple lat-lon grid.

Known issues

The test currently surfaces three Reactant/Oceananigans limitations with distributed TripolarGrid:

  1. set!(field, function) on distributed TripolarGrid fails (unsupported dynamic function invocation in LLVM) — worked around with set_initial_conditions=false
  2. ImmersedBoundaryGrid construction fails for distributed TripolarGrid — worked around with grid_type=:simple_tripolar
  3. Coupled model compilation hits a Reactant MLIR bug (llvm.call op operand type mismatch in raise_triton_custom_call pass) — needs upstream Reactant fix

Test plan

  • CI sharded ocean climate compile job passes
  • CI sharded ocean climate run job passes
  • Existing sharded baroclinic instability tests still pass
  • Existing serial ocean climate tests still pass

🤖 Generated with Claude Code

Add compile and run scripts for sharded ocean climate simulation using a
tripolar grid with the coupled ocean-atmosphere model. This tests sharding
with a different grid topology (tripolar vs lat-lon) and atmosphere coupling.

Changes:
- New sharding/sharded_ocean_climate_simulation_{run,compile}.jl scripts
- Refactor data_free_ocean_climate_model_init to accept (arch, Nx, Ny, Nz)
  for explicit grid sizing needed by sharding tests
- Add grid_type option (:gaussian_islands, :simple_tripolar) and
  set_initial_conditions flag to data_free_ocean_climate_model_init
- Add simple_tripolar_grid helper in model_utils.jl
- Update Compile.yml and Run.yml workflows to include sharded ocean climate

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@glwagner
Copy link
Copy Markdown
Collaborator Author

cc @romanlee @dkytezab @Pangoraw

@Pangoraw
Copy link
Copy Markdown
Collaborator

OOMs again, should we diminish the problem size like in #258 ?

@glwagner
Copy link
Copy Markdown
Collaborator Author

OOMs again, should we diminish the problem size like in #258 ?

I think so. Also baroclinic instability -- hm!

Match the smaller defaults (64×64) used in PR #258 for the sharded
baroclinic instability test to avoid out-of-memory on CI runners.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Switch sharded ocean climate default grid from simple_tripolar to
simple_lat_lon to match the serial ocean climate simulation. Keep
set_initial_conditions=false since set!(field, function) cannot compile
through Reactant's LLVM on any distributed architecture.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants