From 7402408c5fde8a514823b3bb4feca483dfa31cd8 Mon Sep 17 00:00:00 2001
From: Aaron Pham <29749331+aarnphm@users.noreply.github.com>
Date: Fri, 17 Nov 2023 10:40:45 -0500
Subject: [PATCH] fix(envvar): explicitly set NVIDIA_DRIVER_CAPABILITIES (#681)

Signed-off-by: Aaron <29749331+aarnphm@users.noreply.github.com>
---
 openllm-python/src/openllm/bundle/_package.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/openllm-python/src/openllm/bundle/_package.py b/openllm-python/src/openllm/bundle/_package.py
index bed265b8..f885a30c 100644
--- a/openllm-python/src/openllm/bundle/_package.py
+++ b/openllm-python/src/openllm/bundle/_package.py
@@ -98,6 +98,7 @@ def construct_docker_options(
   # XXX: We need to quote this so that the envvar in container recognize as valid json
   environ['OPENLLM_CONFIG'] = f"'{environ['OPENLLM_CONFIG']}'"
   environ.pop('BENTOML_HOME', None)  # NOTE: irrelevant in container
+  environ['NVIDIA_DRIVER_CAPABILITIES'] = 'compute,utility'
   return DockerOptions(
     base_image=oci.RefResolver.construct_base_image(container_registry, container_version_strategy),
     env=environ,