[docker] Change default OMP_NUM_THREADS back to 1 for GPU (deepjavali…

…brary#1073)
Lokiiiiii · Sep 12, 2023 · cece592 · cece592
1 parent 4f79d7a
commit cece592
Show file tree

Hide file tree

Showing 5 changed files with 8 additions and 4 deletions.
diff --git a/engines/python/setup/djl_python/stable_diffusion_inf2.py b/engines/python/setup/djl_python/stable_diffusion_inf2.py
@@ -131,6 +131,10 @@ def __init__(self):
         self.tensor_parallel_degree = None
 
     def initialize(self, properties: dict):
+        # TODO: OMP_NUM_THREADS=1 has sightly negative impact for stable_diffusion model
+        # Remove the following workaround once CLIP model can run on neuron core
+        os.environ.pop("OMP_NUM_THREADS")
+        os.unsetenv("OMP_NUM_THREADS")
         # model_id can point to huggingface model_id or local directory.
         # If option.model_id points to a s3 bucket, we download it and set model_id to the download directory.
         # Otherwise, we assume model artifacts are in the model_dir

diff --git a/serving/docker/deepspeed.Dockerfile b/serving/docker/deepspeed.Dockerfile
@@ -35,7 +35,7 @@ COPY dockerd-entrypoint.sh /usr/local/bin/dockerd-entrypoint.sh
 RUN chmod +x /usr/local/bin/dockerd-entrypoint.sh
 WORKDIR /opt/djl
 ENV JAVA_HOME=/usr/lib/jvm/java-11-openjdk-amd64
-ENV NO_OMP_NUM_THREADS=true
+# ENV NO_OMP_NUM_THREADS=true
 ENV JAVA_OPTS="-Xmx1g -Xms1g -XX:+ExitOnOutOfMemoryError"
 ENV MODEL_SERVER_HOME=/opt/djl
 ENV MODEL_LOADING_TIMEOUT=1200

diff --git a/serving/docker/fastertransformer.Dockerfile b/serving/docker/fastertransformer.Dockerfile
@@ -31,7 +31,7 @@ COPY dockerd-entrypoint.sh /usr/local/bin/dockerd-entrypoint.sh
 RUN chmod +x /usr/local/bin/dockerd-entrypoint.sh
 WORKDIR /opt/djl
 ENV JAVA_HOME=/usr/lib/jvm/java-11-openjdk-amd64
-ENV NO_OMP_NUM_THREADS=true
+# ENV NO_OMP_NUM_THREADS=true
 ENV JAVA_OPTS="-Xmx1g -Xms1g -XX:+ExitOnOutOfMemoryError"
 ENV MODEL_SERVER_HOME=/opt/djl
 ENV MODEL_LOADING_TIMEOUT=1200

diff --git a/serving/docker/pytorch-cu118.Dockerfile b/serving/docker/pytorch-cu118.Dockerfile
@@ -25,7 +25,7 @@ COPY dockerd-entrypoint.sh /usr/local/bin/dockerd-entrypoint.sh
 
 WORKDIR /opt/djl
 ENV JAVA_HOME=/usr/lib/jvm/java-11-openjdk-amd64
-ENV NO_OMP_NUM_THREADS=true
+# ENV NO_OMP_NUM_THREADS=true
 ENV MODEL_SERVER_HOME=/opt/djl
 ENV DJL_CACHE_DIR=/tmp/.djl.ai
 ENV HUGGINGFACE_HUB_CACHE=/tmp

diff --git a/serving/docker/pytorch-inf2.Dockerfile b/serving/docker/pytorch-inf2.Dockerfile
@@ -30,7 +30,7 @@ COPY dockerd-entrypoint.sh /usr/local/bin/dockerd-entrypoint.sh
 RUN chmod +x /usr/local/bin/dockerd-entrypoint.sh
 WORKDIR /opt/djl
 ENV JAVA_HOME=/usr/lib/jvm/java-11-openjdk-amd64
-ENV NO_OMP_NUM_THREADS=true
+# ENV NO_OMP_NUM_THREADS=true
 ENV MODEL_SERVER_HOME=/opt/djl
 ENV DJL_CACHE_DIR=/tmp/.djl.ai
 ENV HUGGINGFACE_HUB_CACHE=/tmp/.cache/huggingface/hub