DB

Braid provenance system

System goals

• Embrace automation in data analysis, retention, decision-making

• Enable users to trace back to how decisions were made

• Necessitates recording what went into model training, including external data, simulations, and structures of other learning and analysis activity

• Envision a versioned database for ML model states with HPC interfaces

• Develop recursive and versioned provenance structures:

  • Models may be constructed via other models (estimates, surrogates)

  • Models are constantly updated (track past decisions and allow updates)

• Integrate with other Braid components

Architecture

Conceptual architecture

Software block diagram

Use cases

SLAC workflow

Goals

1. Perform data reduction at the edge

2. Train models on specific characteristics of experiment-specific data

Workflow

1. Scientist configures experiment parameters

2. Workflow launches simulations with experiment parameters

3. Complete simulations by time experiment data collection is complete

4. Train model on simulation and experiment data

5. Run model on an FPGA to perform data reduction in production

Provenance records

Notes

• Everything has typical metadata like timestamps

• Records can be updated

  • Not immutable history like most provenance data

  • Old versions can be recovered and used

Records

1. Experimental configurations (independent?)

2. Experiment outputs

3. Other simulation inputs?

   1. Software version, configuration?

4. Simulation outputs

5. Training data ingest

6. Inference outputs (statements)

   1. Could be in the form of tests

   2. Like a super-Jenkins

BraggNN workflow

Goals

1. Improve peak finding

2. Train model to represent Bragg peaks

Workflow

1. Scientist configures experiment parameters

2. APS collects raw scattering data

3. Run peak finding on raw data, label peaks

4. Train model on peaks to represent raw data

5. Reproduce and save peak locations

Provenance records

Notes

• Everything has typical metadata like timestamps

Records

1. Experimental configurations (independent?)

2. Experiment outputs

3. Derived peak locations

4. Models trained, checkpoints, etc.

5. Inferred peak locations from trained model

SSX

Goals

1. Track the provenance of SSX crystal structures

Workflow

1. Scientists create a beamline.json and process.phil file

2. Analysis is performed on the input data using these configs to create int files

3. The int files are used with a prime.phil file to create a structure

Provenance records

Notes

1. Structures can come from multiple experiments. This is defined by an intlist in the prime.phil file.

Records

1. Experimental config files (phil, beamline.json)

2. Analysis results (int files)

3. Derived structure

CTSegNet

Goals

1. Track the history of various U-Net-like models used for trial-and-error image segmentation

Workflow

Diagram

1. A tomo scan is obtained

2. Perform image processing, contrast adjustment, etc.

3. Apply (labeling) "masks"

4. Run ensemble models in inference mode

5. Get new segmentations

6. Aggregate segmentation results

7. Re-train models and loop…​

Provenance notes

1. Models are trained on inferences of previous models

2. Provenance queries:

   1. "What data was used to train this model?"

3. Refer to TomoBank IDs for data identification scheme

4. Need rich metadata for search

Samarakoon/Osborn

Goals

1. Fit simulated crystal structure to scattering data

Workflow

1. Obtain neutron scattering data

2. Apply auto-encoder to identify important features

3. Apply dimensionality reduction

4. Fit to data

Getting started

Note: Additional docs forthcoming to help get setup in a (mini)conda based environment.

• Start by installing poetry per https://python-poetry.org/docs/#osx—​linux—​bashonwindows-install-instructions

• Then, run poetry install to setup a local virtual environment from which to run other applications

• Run, poetry run pre-commit install to setup pre-commit hooks for code formatting, lining, etc.

• Tests can be run using scripts in the tests directory.

• Unit tests can be run with the command poetry run pytest pytests/ which will run the pytest test driver from the current virtual environment on all then tests defined in the pytests directory.

Working with FuncX

Using FuncX requires that the funcx-endpoint be installed in a working environment whether it is a conda, pip or otherwise installed process.

Setting up to run with funcx is a multi-step process:

1. Create a new funcX endpoint configured so that it can use the BraidDB library. This can be done with the shell script: scripts/configure_funcx_endpoint.sh. Provide one command line argument: the name of the funcx endpoint to be created/configured. If no name is provided, the endpoint will be named braid_db. This endpoint will be configured such that it has access to the implementation stored in the .venv directory of the braid db project.

2. Start the new endpoint with the command funcx-endpoint start <endpoint_name> where endpoint_name is as configured in the previous step. Take note of the UUID generated for the new endpoint.

3. If you want the funcx hosted braid db to store files in a different location, edit src/braid_db/funcx/funcx_main.py and change the value of DB_FILE in the function funcx_add_record. If not edited the funcx-based operations will store their entries in the file ~/funcx-braid.db.

4. Register the function(s) to be exposed to funcx. This can be done with the command .venv/bin/register-funcx (or poetry run register-funcx). Note that this requires that the command poetry install has previously been run so that the script is installed in the virtual environment (in the .venv/bin directory). As before, take note of the UUID for the registered function.

5. To test invoking the add record function via funcx, run the command: .venv/bin/funcx-add-record --endpoint-id <endpoint_id> --function-id <function_id> using the value for endpoint id and function id output in the previous steps. This should output the record id. One can use a tool like sqlite3 to verify that records are stored in the database file.

Developer notes

• There is a high-level SQL API wrapper in db_tools called BraidSQL.
  This API is generic SQL, it does not know about Braid concepts

• The high-level Braid Database API is called BraidDB

• BraidDB is used by the Braid concepts: BraidFact, BraidRecord, BraidModel, …​

• The Braid concepts are used by the workflows

• We constantly check the DB connection because this is useful when running workflows

Tools

bin/braid-db-create

Creates a DB based on the structure in braid-db.sql

bin/braid-db-print

Print the DB to text

Tests

Tests are in the test/ directory.

Tests are run nightly at:

• https://jenkins-ci.cels.anl.gov/job/Braid-DB-Core

• https://jenkins-ci.cels.anl.gov/job/Braid-DB-Workflows

They are also run via Github Actions for each push or pull request against the origin repo.