Skip to content

Add extra fast srf parser#46

Merged
lispandfound merged 29 commits intomainfrom
rust_srf_parser
Jul 8, 2025
Merged

Add extra fast srf parser#46
lispandfound merged 29 commits intomainfrom
rust_srf_parser

Conversation

@lispandfound
Copy link
Copy Markdown
Contributor

The current SRF parser is fast-ish but not fast enough. Additionally being written in Cython makes it a little annoying to include third party libraries that might make it faster. This PR rewrites the SRF parser in rust and pulls in the lexical_core crate to get 7x speedup in SRF parsing performance with the addition of a more maintainable codebase, type-safety, and memory-safety.

Challenges with Cython

The main challenges with Cython programs are:

  1. Poor 3rd party library support. While you can bundle 3rd party libraries with Cython code (this is just C after all), it isn't easy. In the limit you basically have to bootstrap in a worse version of a C build system alongside your code. The rust equivalent of this is simple. Adding the super fast lexical core library was as simple as adding the line lexical-core = "1.0.5" to the Cargo.toml file in the root of the rust package. The crates registry contains thousands of such packages.
  2. Really bad interpretability. The same line in Cython can have drastically different performance characteristics depending on if the Cython compiler chooses to interpret your code as C or Python. Take for instance the code in the old srf_parser.pyx
    for i in range(matrix.entries):
        # These "=" assignments are done in Python, not C!
        data[i] = matrix.data[i]
        col_indices[i] = matrix.col_ptr[i]
        if max_col_ind <  matrix.col_ptr[i]:
            # This is a regular C assignment
            max_col_ind = matrix.col_ptr[i]

I had to comment the code here to make it clear that one was assigning a value in python (read: slow!) and one was assigning a value in C. This can be a nightmare for trying to write performant extensions because a single line of python in a hot loop can completely trash your performance and the only way to know is with profiling or reading transpiled cython C output. The rust option makes it slightly harder to call Python code but with the huge benefit of making it very clear what is fast rust and what is slow Python.
3. All the problems with C, but worse. Cython is essentially a worse version of C. It retains the all the memory footguns of C but adds weird syntax, odd gotchas, and C89 style variable declarations. Rust is a performant, modern, memory-safe language that makes it easier to write the code once and forget about it.

How much faster did the parser get?

The main issue with the C code is that it spends 99% of its time on lines that look like fscanf("%f", &latitude). C's implementation of fscanf is quite slow, and the somewhat faster strtod requires a lot of ceremony to make work properly. Additionally, the arrays storing all this information unfortunately can't be allocated just once because they don't know their size ahead of time due to quirks of the SRF format. My SRF generation stages were taking an hour to run reading and writing SRFs multiple times over in the intermediate stages. This would essentially mean days of turn around time for alpine fault simulations. To make this faster, the rust version of this code uses the incredibly fast lexical core crate. This crate uses SIMD instructions to process multiple characters in parallel which provides an enormous speedup. How big? Well the old code reads my 18GB 100m alpine_hope_1.srf file in nearly 5 minutes and the new code takes just 30 seconds

(fast) jfa92@hypocentre:~/source_modelling_comp$ /usr/bin/time -v python test.py
	Command being timed: "python test.py"
	User time (seconds): 34.18
	System time (seconds): 5.41
	Percent of CPU this job got: 108%
	Elapsed (wall clock) time (h:mm:ss or m:ss): 0:36.33
	Average shared text size (kbytes): 0
	Average unshared data size (kbytes): 0
	Average stack size (kbytes): 0
	Average total size (kbytes): 0
	Maximum resident set size (kbytes): 62311596
	Average resident set size (kbytes): 0
	Major (requiring I/O) page faults: 0
	Minor (reclaiming a frame) page faults: 362633
	Voluntary context switches: 54
	Involuntary context switches: 152
	Swaps: 0
	File system inputs: 0
	File system outputs: 0
	Socket messages sent: 0
	Socket messages received: 0
	Signals delivered: 0
	Page size (bytes): 4096
	Exit status: 0

The effect is more pronounced on newer machines with larger caches and better SIMD instruction sets, my home desktop has a two-year-old 12th gen mid range i5 and reads a 5gb SRF in only 8 seconds.

Structure of a Rust binding

I intend this PR to be a demo of how to write such rust bindings. It's at least as easy as Cython to setup, and much easier if you want to add a third-party library. The steps I went through are as follows:

  1. Navigate to the top-level of the python package you want to install a rust component for. In this case, the source_modelling directory.
~/source_modelling $ cd source_modelling
~/source_modelling/source_modelling $ ls
-rw-r--r--  1 jake jake  67K Feb 18 23:21 ccldpy.py
-rw-r--r--  1 jake jake  16K Apr  8 00:01 community_fault_model.py
-rw-r--r--  1 jake jake  20K Jun  4 04:35 fsp.py
-rw-r--r--  1 jake jake 5.9K May  1 02:22 gsf.py
[...]
  1. Create a new directory for your rust module. I chose srf_parser. Change into this directory and run cargo init --lib. The --lib flag tells cargo to make a library project rather than an executable (the default). If you don't have rust installed you can very easily do so with rustup. I have globally installed rust on Hypocentre for researchers. It is also bundled with popular distributions like Ubuntu because many of their own tools are written in rust these days.
~/source_modelling/source_modelling $ mkdir srf_parser
~/source_modelling/source_modelling/srf_parser $ cargo init
~/source_modelling/source_modelling/srf_parser $ tree
.
├── Cargo.toml
└── src
    └── lib.rs
  1. Edit the Cargo.toml to compile a Python-compatible library extension. This is as simple as changing the crate type. Add a line like this
[lib]
crate-type = ["cdylib"]

and you're done!
4. Add the pyo3 and numpy crates. You can do this with cargo add. Edit the Cargo.toml to add the extension-module feature.

~/source_modelling/source_modelling/srf_parser $ cargo add pyo3 numpy

and in the cargo file

[dependencies]
# What cargo will add
# pyo3 = "0.25.1"
# Change the above to
pyo3 = { version = "0.25.1", features = ["extension-module"] }
numpy = "0.25"
  1. Open src/lib.rs and start writing your rust code! This part requires knowing rust, which can be a bit of a hurdle but it isn't that hard to learn (no harder than learning to write Cython imo). The only things different you need to know is pyo3 specific stuff. To make a rust function callable from python you add #[pyfunction] to the top and add the py argument which is a copy of the python interpreter:
#[pyfunction]
fn parse_srf(
    py: Python<'_>,
    file_path: &str,
    offset: usize,
    num_points: usize,
) -> PyResult<(Py<PyAny>, Py<PyAny>)> {

Then you can add it to your module at the bottom

rustrust#[pymodule]
fn srf_parser(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(parse_srf, m)?)?;
    Ok(())
}
  1. Now that you're ready to bundle your shiny rust code with your python project you can add it to the pyproject.toml to bundle everything together. We use setuptools mainly so this is a touch more challenging than normal but it is still pretty simple. You just add setuptools-rust to the build system requirements
[build-system]
requires = ["setuptools", "setuptools-scm",  "setuptools-rust>=1.1.0"]
build-backend = "setuptools.build_meta"

and then add a line to the setup.py to instruct python to build the moment when you pip install

from setuptools import setup
from setuptools_rust import RustExtension, Binding
import os

setup(
    name="source_modelling",
    packages=["source_modelling"],
    rust_extensions=[RustExtension("source_modelling.srf_parser", "source_modelling/srf_parser/Cargo.toml", binding=Binding.PyO3, native=True)],
    include_package_data=True,
    use_scm_version=True,
    zip_safe=False,  # required for Rust extensions
)

I added native=True to compile the module with the fastest possible native instruction set when the module is built. Generally you shouldn't do this I just did it to make sure lexical_core gets to use SIMD instructions.
7. Pip install and go! You can import your module in python like an ordinary one

from source_modelling import srf_parser
srf_parser.parse_srf(...)

Now you can distribute this to researchers. Like a cython extension, they don't need to know anything about the distribution. If the wheel is prebuilt they don't even need to have rust installed. Unfortunately we use git packages so they need rust to compile this wheel themselves. I've installed rust on Hypocentre with an automatic update option turned on so they don't have to think about it there.

Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Summary of Changes

Hello @lispandfound, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a significant performance improvement to the SRF file parsing by rewriting the existing Cython-based parser in Rust. The new Rust parser leverages highly optimized libraries for numerical parsing, resulting in a substantial speedup and a more robust, maintainable, and memory-safe codebase. It also serves as a demonstration of integrating Rust extensions into a Python project.

Highlights

  • Performance Enhancement: I've replaced the existing Cython SRF parser with a new Rust implementation, achieving a significant speedup in SRF file parsing. For instance, an 18GB SRF file that previously took nearly 5 minutes to parse now completes in approximately 30 seconds, representing up to a 7x performance improvement.
  • Technology Migration & Code Quality: The core SRF parsing logic has been migrated from Cython to Rust. This transition addresses Cython's limitations regarding third-party library integration and performance interpretability, resulting in a more maintainable, type-safe, and memory-safe codebase.
  • Simplified SRF Data Model: The new parser primarily focuses on the slip1 component. This simplifies the internal representation of SRF slip data by effectively deprecating explicit slip2 and slip3 components within the SrfFile object and streamlining related parsing and writing logic.
  • Rust-Python Integration: This pull request serves as a practical demonstration of integrating Rust extensions into a Python project. It showcases the setup and build process using pyo3 for Python bindings and setuptools-rust for seamless distribution, including dependency management via Cargo.toml.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in issue comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request replaces the Cython-based SRF parser with a new implementation in Rust, achieving a significant performance boost. The use of mmap and lexical-core is a great choice for this task. The code is well-structured and the PR description is very clear and helpful. The review includes suggestions to further improve performance, robustness, and maintainability.

@lispandfound
Copy link
Copy Markdown
Contributor Author

Made some refactors suggested by Gemini (although what it suggested made things even worse so I had to undo them and write my own solution). Basically there were too many copies floating around in my code. Now the copying has been reduced I can parse the 18GB alpine SRF on my own desktop in about 40s. Can't test on Hypocentre right now because it is being hammered atm but I imagine it also faster than before with better memory usage.

@lispandfound lispandfound requested a review from Copilot July 6, 2025 23:41

This comment was marked as outdated.

claudio525
claudio525 previously approved these changes Jul 6, 2025
@lispandfound
Copy link
Copy Markdown
Contributor Author

lispandfound commented Jul 7, 2025

While waiting for a velocity model to generate I tinkered with this a bit more and disabling 64-bit indexing on the SRF array because I figure million second long SRF indexing windows ought to be long enough for anyone. This brings the memory usage down to the optimal amount (basically the size of the SRF in ascii plus the exact memory usage required for the python arrays). If we ever have an SRF that wraps around the whole world and this becomes an issue then we have other problems.

UPDATE: doesn't matter because scipy will "helpfully" upgrade all the arrays in question to 64-bit (signed!) integers, forcing a copy to occur and more memory usage. The best I can do is match the scipy integer format and use i64 instead of usize so the csr array in scipy doesn't reinterpret the data. Because of the aforementioned copy we get to do this "for free" since we were paying the 64-bit tax regardless.

claudio525
claudio525 previously approved these changes Jul 7, 2025
Copy link
Copy Markdown
Contributor

@AndrewRidden-Harper AndrewRidden-Harper left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

very nice

@lispandfound
Copy link
Copy Markdown
Contributor Author

Sorry @claudio525 one further change, the SRF writer now also uses rust. Writing takes just about as much time as reading (roughly 40 seconds on hypocentre), which suggests that the bottleneck is probably still (de)serialising a lot of floats. The SRFs produced by write_srf tend to be smaller than genslip ones.

claudio525
claudio525 previously approved these changes Jul 7, 2025
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull Request Overview

This PR replaces the existing Cython-based SRF parser with a Rust implementation, integrates it into the Python package, and updates tests and build configuration accordingly.

  • Remove srf_reader.pyx and Cython build, add Rust-based srf_parser crate for faster parsing.
  • Update read_srf/write_srf in srf.py to call the new Rust functions and drop multi-component slip support.
  • Adjust tests to reflect only a single slip component and added a windowed slip-value check.

Reviewed Changes

Copilot reviewed 7 out of 8 changed files in this pull request and generated 4 comments.

Show a summary per file
File Description
tests/test_srf.py Modified tests to drop slip2/3 assertions and added end-window check
source_modelling/srf_reader.pyx Removed old Cython-based SRF reader entirely
source_modelling/srf_parser/src/lib.rs Introduced Rust parser (parse_srf, write_srf_points, etc.)
source_modelling/srf_parser/Cargo.toml Defined the Rust crate with dependencies
source_modelling/srf.py Switched to Rust parser calls, removed slip2/3 handling
setup.py Replaced Cython extension with setuptools-rust extension
pyproject.toml Updated build-system requirements for Rust
Comments suppressed due to low confidence (2)

source_modelling/srf_parser/src/lib.rs:112

  • [nitpick] The variable _nt2 actually holds the slip2 float value (and _slip2 holds the count); consider renaming them to slip2_val and nt2 respectively for clarity.
        let _nt2 = parse_value::<f32>(data, &mut index)?;

source_modelling/srf.py:442

  • There are currently no tests targeting the new write_srf_points or write_srf functionality; consider adding tests to verify the output SRF format and data correctness.
    srf_parser.write_srf_points(

@lispandfound lispandfound merged commit ee4ab7d into main Jul 8, 2025
6 checks passed
@lispandfound lispandfound deleted the rust_srf_parser branch July 8, 2025 02:39
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.

4 participants