diff --git a/examples/cv_kot_example/nvidia_image_test.py b/examples/cv_kot_example/nvidia_image_test.py new file mode 100644 index 00000000..29d0026d --- /dev/null +++ b/examples/cv_kot_example/nvidia_image_test.py @@ -0,0 +1,13 @@ +import ebonite +from ebonite.build.builder.base import use_local_installation +from ebonite.ext.cv_streaming.runner import CvDockerRunner +from ebonite.ext.cv_streaming.server import CvStreamingServer + +with use_local_installation(): + ebnt = ebonite.client.Ebonite.local(clear=True) + task = ebnt.get_or_create_task('Proj', 'Task') + model = ebnt.create_model(lambda x: x, 'test_input', 'model', task_name='Task') + env = ebnt.get_default_environment() + env.Params.default_runner = CvDockerRunner() + image = ebnt.create_image(model, 'image_name', server=CvStreamingServer, builder_args={'force_overwrite': True}) + instance = ebnt.create_instance(image, 'instance', runner_kwargs={'gpus': 1}, run=True) diff --git a/src/ebonite/ext/cv_streaming/__init__.py b/src/ebonite/ext/cv_streaming/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/ebonite/ext/cv_streaming/build/app.py b/src/ebonite/ext/cv_streaming/build/app.py new file mode 100644 index 00000000..5c928756 --- /dev/null +++ b/src/ebonite/ext/cv_streaming/build/app.py @@ -0,0 +1,6 @@ +from ebonite import start_runtime +from ebonite.ext.cv_streaming.server import CvStreamingServer + + +# STUB APP +start_runtime(CvStreamingServer) diff --git a/src/ebonite/ext/cv_streaming/runner.py b/src/ebonite/ext/cv_streaming/runner.py new file mode 100644 index 00000000..5fecd672 --- /dev/null +++ b/src/ebonite/ext/cv_streaming/runner.py @@ -0,0 +1,28 @@ +import subprocess + +from ebonite.ext.docker.base import DockerContainer, DockerEnv, DockerImage +from ebonite.ext.docker.runner import DockerRunner + + +class CvDockerRunner(DockerRunner): + + def run(self, instance: DockerContainer, image: DockerImage, env: DockerEnv, rm=True, detach=True, **kwargs): + if not (isinstance(instance, DockerContainer) and isinstance(image, DockerImage) and + isinstance(env, DockerEnv)): + raise TypeError('DockerRunner works with DockerContainer, DockerImage and DockerHost only') + + gpus = kwargs.get('gpus') + if gpus is None: + gpus = 'all ' + + ports = '' + for exp, pub in instance.port_mapping.items(): + ports += f'-p {exp}:{pub} ' + + if rm: + rm = '--rm ' + else: + rm = '' + + command = f"docker run --gpus {gpus} {rm} {ports} --name {instance.name} -d {image.uri}" + subprocess.run(command.split()) diff --git a/src/ebonite/ext/cv_streaming/server.py b/src/ebonite/ext/cv_streaming/server.py new file mode 100644 index 00000000..f9be55d6 --- /dev/null +++ b/src/ebonite/ext/cv_streaming/server.py @@ -0,0 +1,19 @@ +from ebonite.runtime.interface.base import Interface +from ebonite.runtime.server.base import Server +from ebonite.utils.fs import current_module_path + + +class CvStreamingServer(Server): + + additional_sources = [ + current_module_path('build', 'app.py') # replace stub in base image + ] + + additional_options = {'docker': { + 'templates_dir': current_module_path('build'), + 'run_cmd': False, # base image has already specified command + 'base_image': 'nvcr.io/nvidia/tensorrt:20.07.1-py3' + }} + + def run(self, executor: Interface): + pass diff --git a/tests/ext/test_cvstreaming/__init__.py b/tests/ext/test_cvstreaming/__init__.py new file mode 100644 index 00000000..e69de29b