My foray into BMI: bmi-operator, a Kubernetes-based orchestration tool

Hello folks - I learned about the BMI ecosystem a few months ago, and I was inspired to put together and share the project I’ve worked on. Thanks to @mdpiper for a friendly conversation and encouragement!

I am calling it the bmi-operator (justin/bmi-operator, personal webhost link). It’s an orchestration tool that takes BMI models and runs them on Kubernetes, the distributed computing platform. Its main feature is queueing and durability: Kubernetes orchestrates the simulations and stores result data, so it’s easy to start many simulations with slight variations and check in on them, and access the results later.

It does also allow you to run multi-model scenarios: similar to how Landlab supports multiple components operating on a shared grid - but distributed! I did replicate a few Landlab tutorials using the tool to test it out, like the transport-length hillslope diffuser tutorial (my copy). I also included a version of the cem-grpc4bmi tutorial (my version). The interface looks roughly like this:

spec = bmi_operator.run.Spec(
    name="tlhdiff-tutorial",
    models=[flow_director, diffuser],  # BMI model container parameters set up earlier
    dt="1000yr",
    end_time="2000000yr",
    outputs_domain=outputs,            # uses cloud storage for results
    variables=[
        bmi_operator.run.Variable(     # set initial data and continuously save the variable through the simulation
            "topographic__elevation",
            seed=starting_data["topo.npy"],
            every="100kyr",
        )
    ],
)

run = client.runs.submit(spec)
run.wait(timeout=1800)

results = run.results()
times, elevations = results.topographic__elevation.every("100kyr")

And for the evidence, here’s a visualization you might make of the resulting timeseries data:
Animated evolution of the TLHDiff sloped landscape

As for how it’s built: it uses the Kubernetes operator pattern, relies on the BMI models being containerized and integrated with grpc4bmi, and uses helm for installation. There’s a Python client library that is meant to be used in a Jupyter notebook to interact with it. If you have Docker, you can try it out! For someone developing a new model, I don’t think it would be useful. But it has some capabilities for someone who would be trying to use existing models, running many variations or long-running tasks.

I’ve enjoyed working on this and exploring some of the modeling work and research being done using BMI. I’m an infrastructure software engineer by trade, so, I thought an infrastructure project using tools I know would be a good way to take a first look at some of what CSDMS has built - and sometimes the best payoff of good API design, in my experience, is the platforms and tools you can easily build on top of it!

I’m happy to discuss the project, but even moreso, would love to contribute to the open-source BMI and CSDMS ecosystem where I can. Feel free to share your thoughts or email me at justin@palpant.us, I’m happy to chat, receive feedback, and learn!

Honest disclosure: I also used this project to gain some more experience with coding using LLMs - it’s taking off at my work, and I wanted to practice. I have kept a close eye and done a fair amount of coding the normal way, but certainly there are signs, styles, and likely, errors, from trying to develop in this way. It’s very alpha-stage work, and I’m also happy to share some of the pitfalls I ran into.

2 Likes

It’ll take me a bit to fully understand your work, but it looks great! Thank you for sharing this!

Great to see usage of grpc4bmi. This implements my long outstanding issue at Kubernetes support · Issue #27 · eWaterCycle/grpc4bmi · GitHub .

How are you handling the config file, input files and output files? Is it not needed for your use case? Or are you using some mounted filesystem?

It would be really nice if improvements in forks could be merged back to GitHub - eWaterCycle/grpc4bmi: gRPC wrapper for model with a Basic modeling interface · GitHub

Amazing, I hadn’t seen that there has been prior discussion about Kubernetes! And I would be very happy to upstream some of the changes to grpc4bmi that I have made, I will make an effort to do so soon.

This initial version has support for a single config file, like one whose filepath would be passed to initialize(configFilePath). The contents of that file could be set for each model, for each run:

cem_model = bmi.run.Model(
    name="cem",
    image="registry.palpant.us/justin/cem-grpc4bmi:latest",
    env={"BMI_MODEL": "cem"},
    config="""50, 1000, 1000.0, 1
0.01, 10.0, 0.001
""",
)

config makes sure the contents are available to the model container as a file, and then the file name is put in the call to initialize(); some models prefer to get the file text in the initialize call; for those, bmi.run.Model() takes a configString argument (mutually exclusive with config).

Behind the scenes, config generates a Kubernetes ConfigMap from the contents, mounts that to a known path, and then callsinitialize(mountedFilePath) when the simulation starts.

However, I haven’t added any method of supporting inputs and outputs that take the form of files - only inputs and outputs via BMI variables, using get_value and set_value, for variables whose names are in get_input_var_names and get_output_var_names. It is a serious limitation: I could not use models which require more complex inputs via files, or which produce PNGs or other file outputs. Since the BMI spec doesn’t define input or output files, I’d like to find a generic way to support this, but am not sure yet how.

I would like to think about it more: would you happen to have an example BMI model container that uses input and output files that I could test on?

The most simple eWaterCycle model is in the GitHub - eWaterCycle/leakybucket-bmi: Template project for a hydrological model BMI for eWaterCycle. · GitHub repo with Docker image ghcr.io/ewatercycle/leakybucket-grpc4bmi:v0.0.1. It needs a json config file like

{
 "leakiness": 0.05,
 "precipitation_file": "pr.nc",
 "temperature_file": "tas.nc"
}

Where pr.nc and tas.nc are paths netcdf files with forcings.

These forcings can be generated with eWaterCycle package see notebook in the repo.

1 Like

Thank you for the example! I have been playing around with things, and I think the path I will try to support would be to upload the files to storage, and then arrange for the operator to download them to the model Pod. There is already some facility like this for describing storage and setting initial conditions, like setting topographic__elevation with a DEM, and also for storing results. This is inconvenient if the input is extremely large and needed by multiple models, however. But in a situation where the client and model are not running in the cluster, cloud storage for the exchange of binary data is very universal.

If I finish this, the API would probably resemble:

forcing = client.assets.upload(
    {
        "pr.nc": cmip_forcing["pr"],
        "tas.nc": cmip_forcing["tas"],
    },
    domain=inputs,
)

mounted = model.attach(forcing, name="forcing") # <- this is not built yet

config = {
    "leakiness": 0.05,
    "precipitation_file": mounted["pr.nc"], # the returned object says where the model Pod should expect to find the file content
    "temperature_file": mounted["tas.nc"],
}

model.config = json.dumps(config) # this already exists and works today!

What do you think? Would automatic exchange via cloud storage be reasonable for most cases?

That looks like a fine API. I am assuming cmip_forcing["pr"] would be some S3 public url or presigned url.

Using mounted cloud storage is useful for models that require posix files. For more modern models then zarr or opendap urls in config will work so you do not need upload/attach for those

Some models could have many input files like for example LISFLOOD has LISFLOOD static and parameter maps for Europe dataset . For those it would be nice to upload/attach a whole nested directory.