https://github.com/ogvalt/triton-testcontainer
Package for running Nvidia Triton within python test with features like Dockerfile DSL and building images on fly.
https://github.com/ogvalt/triton-testcontainer
docker dockerfile pytest python testcontainers triton
Last synced: 11 months ago
JSON representation
Package for running Nvidia Triton within python test with features like Dockerfile DSL and building images on fly.
- Host: GitHub
- URL: https://github.com/ogvalt/triton-testcontainer
- Owner: ogvalt
- License: apache-2.0
- Created: 2024-04-03T15:46:45.000Z (about 2 years ago)
- Default Branch: main
- Last Pushed: 2024-05-21T15:16:52.000Z (about 2 years ago)
- Last Synced: 2025-04-19T11:52:04.831Z (about 1 year ago)
- Topics: docker, dockerfile, pytest, python, testcontainers, triton
- Language: Python
- Homepage:
- Size: 367 KB
- Stars: 4
- Watchers: 1
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# Triton Testcontainer
This package started as an implementation of testcontainer module for Nvidi Triton, but also added other functionality that extends beyond simply running containers withing pytest.
## Usage
This package provides several functions:
* triton container (class `TritonContainer`): manage tritonserver container within a testsuite.
* dockerfile builder (class `DockerfileBuilder`): Dockerfile DSL for generating dockerfiles programatically.
* docker image builder (class `ImageBuilder`): Builing images on fly, e.g. within a testsuite.
## Installation
```bash
pip3 install git+https://github.com/ogvalt/triton-testcontainer.git@v0.6.1
```
## Examples
```python
import pytest
import triton_testcontainer as tritoncontainer
@pytest.fixture(scope="session") # Specify scope of fixture
def setup_triton(request):
# start triton
triton = tritoncontainer.TritonContainer() # Create container instance
triton.start() # launches container, waits until it's ready
def remove_triton_container():
triton.stop() # stops container
request.addfinalizer(remove_triton_container) # handles container at the end of testing session
return triton.get_client() # returns triton http client
def test_example(setup_triton):
triton_client: 'tritonclient.http.InferenceServerClient' = setup_triton
assert True
def test_example_two():
cmd = tritoncontainer.TritonCommand(model_repository=["/models"]).build() # command to run tritonserver with
maps = [{"host": "/path/to/repository", "container": "/models"}] # map repository on host to container
# use context manager to run container on __enter__ and stop it on __exit__
with tritoncontainer.TritonContainer(with_gpus=True, volume_mapping=maps, command=cmd) as service:
triton_client = service.get_client()
assert True
```
```python
from testcontainers.core.container import DockerContainer
from testcontainers.core.waiting_utils import wait_for_logs
from triton_testcontainer import ImageBuilder, DockerfileBuilder
def test_example_three():
dockerfile = (DockerfileBuilder() \
.from_("ubuntu:20.04") \
.cmd("echo", "hello world")
.build())
with ImageBuilder("hello-world").from_string(context=".", string_dockerfile=dockerfile).ctx_manager() as image:
with DockerContainer(image.tags[0]) as container:
delay = wait_for_logs(container, "hello world", 30)