tox.ini

[tox]
envlist = lint, gen-configs, run-benches

[testenv:gen-configs]
description = generate sample configurations for all plugins
deps =
    ruamel.yaml>=0.18
skip_install = true
commands = 
    python scripts/generate_sample_configurations.py {posargs:sample-configurations}

[testenv:verify-configs]
description = verify that sample configurations for all plugins are properly generated
skip_install = true
commands = 
    bash scripts/verify_generated_configurations.sh {posargs:sample-configurations}
allowlist_externals = bash

# put this here first, consider moving it later
[testenv:run-benches]
description = run benchmarks 
skip_install = true
deps = 
    matplotlib # this is for plotting benchmark comparisons in compare_with_reference.py
    packaging # this is required for flash-attn dep as fms_hf_tuning did not specify
    -e {toxinidir}/plugins/framework # install the framework here as the flash attention deps requires torch
passenv = * # will pass the parent env, otherwise there are too many envs e.g. TRANSFORMERS that need to be set
commands = 
    # need a version of fms-hf-tuning that has integrated the framework
    # NOTE: have to install this first coz havnt merged
    # - this repo has a lot of pins, so we just install it first
    pip install "fms-hf-tuning @ git+https://github.com/foundation-model-stack/fms-hf-tuning.git@"{env:FHT_BRANCH:main}

    # some models need this for tokenizers
    pip install protobuf

    # install the plugins for test
    # NOTE: when there are more plugins install here
    python -m fms_acceleration.cli install -e {toxinidir}/plugins/accelerated-peft
    python -m fms_acceleration.cli install -e {toxinidir}/plugins/fused-ops-and-kernels
    python -m fms_acceleration.cli install -e {toxinidir}/plugins/attention-and-distributed-packing
    python -m fms_acceleration.cli install -e {toxinidir}/plugins/accelerated-moe

    # install the flash attn at the last 
    pip install flash-attn

    # run the benchmark script
    bash scripts/run_benchmarks.sh {posargs:"1 2" "4 8" benchmark_outputs}

allowlist_externals = bash

[testenv:lint]
description = run linters
skip_install = true
deps =
    black>=22.12
    isort>=5.11
commands = 
    black scripts
    isort scripts