To contribute to this project, you need the following:
- UV
- pre-commit (via
uv tool install pre-commit
) - ruff (via
uv tool install ruff
)
- Go 1.22
- Mage (via
go install github.com/magefile/mage@v1.15.0
) - protoc-gen-go (via
go install google.golang.org/protobuf/cmd/protoc-gen-go@v1.33.0
) - Mockery (via
go install github.com/vektra/mockery/v2@v2.43.2
) - Golangci-lint (via
go install github.com/golangci/golangci-lint/cmd/golangci-lint@v1.59.1
)
See Protocol Buffer Compiler Installation
Alternatively, you can use the development container that includes all the required tools.
to configure all the development environment just run mage
target:
mage configure
it will configure MLflow and all the Python dependencies required by the project or run each step manually:
# Install our Python package and its dependencies
pip install -e .
# Install the dreaded psycho
pip install psycopg2-binary
# Archive the MLflow pre-built UI
tar -C /usr/local/python/current/lib/python3.8/site-packages/mlflow -czvf ./ui.tgz ./server/js/build
# Clone the MLflow repo
git clone https://github.com/jgiannuzzi/mlflow.git -b master .mlflow.repo
# Add the UI back to it
tar -C .mlflow.repo/mlflow -xzvf ./ui.tgz
# Install it in editable mode
pip install -e .mlflow.repo
To start the mlflow-go dev server connecting to postgres just run next mage
target:
mage dev
The postgres database should already be running prior to this command. By default service uses next connection string:
postgresql://postgres:postgres@localhost:5432/postgres
but it could be configured in mage
If you wish to contribute to the porting of an existing Python endpoint, you can read our dedicated guide.
The Python integration tests have been adapted to also run against the Go implementation.
Next mage
targets are available to run different types of tests:
# Run all the available tests
mage test:all
# Run just MLflow Python tests
mage test:python
# Run specific MLflow Python tests (matches all tests containing the argument)
mage test:pythonSpecific <test_file::test_name>
#Example
mage test:pythonSpecific ".mlflow.repo/tests/tracking/test_rest_tracking.py::test_rename_experiment"
# Run just unit tests
mage test:unit
Additionally, there is always an option to run, specific test\tests if it is necessary:
pytest tests/tracking/test_rest_tracking.py
To run only the tests targeting the Go implementation, you can use the -k
flag:
pytest tests/tracking/test_rest_tracking.py -k '[go-'
If you'd like to run a specific test and see its output 'live', you can use the -s
flag:
pytest -s "tests/tracking/test_rest_tracking.py::test_create_experiment_validation[go-postgresql]"
See the pytest documentation for more details.
# Build the Go binary in a temporary directory
libpath=$(mktemp -d)
python -m mlflow_go.lib . $libpath
# Run the tests (currently just the server ones)
MLFLOW_GO_LIBRARY_PATH=$libpath pytest --confcutdir=. \
.mlflow.repo/tests/tracking/test_rest_tracking.py \
.mlflow.repo/tests/tracking/test_model_registry.py \
.mlflow.repo/tests/store/tracking/test_sqlalchemy_store.py \
.mlflow.repo/tests/store/model_registry/test_sqlalchemy_store.py \
-k 'not [file'
# Remove the Go binary
rm -rf $libpath
# If you want to run a specific test with more verbosity
# -s for live output
# --log-level=debug for more verbosity (passed down to the Go server/stores)
MLFLOW_GO_LIBRARY_PATH=$libpath pytest --confcutdir=. \
.mlflow.repo/tests/tracking/test_rest_tracking.py::test_create_experiment_validation \
-k 'not [file' \
-s --log-level=debug