chore: Add triton_distributed_rs wheel install to container build

[rmccorm4]

Changes:

• Fix triton_distributed_rs python binding install when using --mount-workspace by using maturin build instead of maturin develop - see refactor: Modify rust vllm example to use container #132 (comment)
  • current code without this PR works with ./container/run.sh -it, but fails with addition of --mount-workspace
  • this PR fixes support for both with or without --mount-workspace
• Use release build of crate instead of debug build with -r in maturin build -r
• Optimize unnecessary docker rebuilds from COPY commands
• Add sanity test to check binding is installed in container build

Closes #148

[rmccorm4]

Moved triton_distributed_rs install outside of venv so it can be more easily tested in current CI setup.

Passes locally too:

root@ced35d0-lcedt:/workspace# pytest -m "mypy or pre_merge" runtime/rust/python-wheel/
/usr/local/lib/python3.12/dist-packages/pytest_asyncio/plugin.py:207: PytestDeprecationWarning: The configuration option "asyncio_default_fixture_loop_scope" is unset.
The event loop scope for asynchronous fixtures will default to the fixture caching scope. Future versions of pytest-asyncio will default the loop scope for asynchronous fixtures to function scope. Set the default fixture loop scope explicitly in order to avoid unexpected behavior in the future. Valid fixture loop scopes are: "function", "class", "module", "package", "session"

  warnings.warn(PytestDeprecationWarning(_DEFAULT_FIXTURE_LOOP_SCOPE_UNSET))
=================================================================== test session starts ===================================================================
platform linux -- Python 3.12.3, pytest-8.3.4, pluggy-1.5.0
benchmark: 5.1.0 (defaults: timer=time.perf_counter disable_gc=False min_rounds=5 min_time=0.000005 max_time=1.0 calibration_precision=10 warmup=False warmup_iterations=100000)
rootdir: /workspace
configfile: pyproject.toml
plugins: mock-3.14.0, timeout-2.3.1, cov-6.0.0, asyncio-0.25.3, benchmark-5.1.0, md-report-0.6.3, mypy-0.10.3, anyio-4.8.0
asyncio: mode=Mode.AUTO, asyncio_default_fixture_loop_scope=None
collected 21 items

runtime/rust/python-wheel/examples/bls/bar.py ..                                                                                                    [  9%]
runtime/rust/python-wheel/examples/bls/bls.py .                                                                                                     [ 14%]
runtime/rust/python-wheel/examples/bls/foo.py .                                                                                                     [ 19%]
runtime/rust/python-wheel/examples/error_handling/__init__.py .                                                                                     [ 23%]
runtime/rust/python-wheel/examples/error_handling/client.py .                                                                                       [ 28%]
runtime/rust/python-wheel/examples/error_handling/run.py .                                                                                          [ 33%]
runtime/rust/python-wheel/examples/error_handling/server.py .                                                                                       [ 38%]
runtime/rust/python-wheel/examples/hello_world/client.py .                                                                                          [ 42%]
runtime/rust/python-wheel/examples/hello_world/run.py .                                                                                             [ 47%]
runtime/rust/python-wheel/examples/hello_world/server.py .                                                                                          [ 52%]
runtime/rust/python-wheel/examples/pipeline/backend.py .                                                                                            [ 57%]
runtime/rust/python-wheel/examples/pipeline/frontend.py .                                                                                           [ 61%]
runtime/rust/python-wheel/examples/pipeline/middle.py .                                                                                             [ 66%]
runtime/rust/python-wheel/examples/pipeline/pipeline.py .                                                                                           [ 71%]
runtime/rust/python-wheel/examples/typed/__init__.py .                                                                                              [ 76%]
runtime/rust/python-wheel/examples/typed/client.py .                                                                                                [ 80%]
runtime/rust/python-wheel/examples/typed/protocol.py .                                                                                              [ 85%]
runtime/rust/python-wheel/examples/typed/server.py .                                                                                                [ 90%]
runtime/rust/python-wheel/python/tests/test_bindings_install.py ..                                                                                  [100%]
========================================================================== mypy ===========================================================================

Success: no issues found in 19 source files
=================================================================== 21 passed in 0.38s ====================================================================

[nnshah1]

LGTM - could move to virtual env - could do in seperate MR

[rmccorm4]

Need to reconcile:

1. pytest tests not using venv in ci job
2. vllm disagg example using venv due to vllm install issue, but not getting the system interpreter packages (triton_distributed_rs) made available to the venv

Need to make both test and vllm example work

[rmccorm4]

triton_distributed_rs package now currently built once, but installed both inside and outside of venv to cover both cases

1. pytest gitlab ci job - currently doesn't activate virtual environment
2. vllm disagg rust example - currently based on virtual environment

./container/build.sh
./container/run.sh -it     # or --mount-workspace for local dev/debug

$ python3 -c "import triton_distributed_rs"  # OK
$ source /opt/triton/venv/bin/activate
(venv) $ python3 -c "import triton_distributed_rs"   # OK

In future we may shift to using virtual environment everywhere and remove the extra pip install to system environment, including in CI test scenarios - but it's not currently there.

Also note that uv pip install --system won't work without something like env var PIP_BREAK_SYSTEM_PACKAGES=1, hence the separate pip install outside of venv in current solution.

[GuanLuo]

I think if you do maturin build instead of maturin develop, then virtual environment is not required

[rmccorm4]

I did have that originally f81c240, but the issue with having everything installed in system environment (not the virtual environment) - is that there were unresolved issues with trying to pip install vllm for the examples/python_rs/llm/vllm example in the system environment. See #134.

So for now, to keep both old pure python examples/tests working, and new mixed python/rust examples working - I went for this approach.

I'd prefer to move towards all or nothing in terms of virtualenv usage or not - but I think we'd need to figure out the issues with the vllm install to remove the virtualenv.

I decided to move forward with test/pipeline improvements before going back to debug the vllm install issue, because I figured we'd make a decision in the near future to install vllm>=0.7 in the container at some point when doing ./build.sh --framework vllm anyways - but may need to avoid breaking the pure-python vllm disagg example (which patches vllm install) when doing so.