Set logging env variable for djl-serving

Author: siddvenk

src/sagemaker/djl_inference/defaults.py

@@ -38,7 +38,7 @@
 }
 
 ALLOWED_INSTANCE_FAMILIES = {
-    "ml.g4",
+    "ml.g4dn",
     "ml.g5",
     "ml.p3",
     "ml.p4",

src/sagemaker/djl_inference/model.py

@@ -14,6 +14,7 @@
 from __future__ import absolute_import
 
 import json
+import logging
 import os.path
 from enum import Enum
 from typing import Optional, Union, Dict
@@ -29,6 +30,19 @@
 from sagemaker.workflow.entities import PipelineVariable
 
 
+logger = logging.getLogger("sagemaker")
+
+
+LOG_LEVEL_MAP = {
+    logging.INFO: "info",
+    logging.DEBUG: "debug",
+    logging.WARNING: "warn",
+    logging.ERROR: "error",
+    logging.FATAL: "fatal",
+    logging.CRITICAL: "fatal",
+    logging.NOTSET: "off",
+}
+
 class DJLEngine(Enum):
     DEEPSPEED = ("DeepSpeed", "djl_python.deepspeed")
     FASTER_TRANSFORMERS = ("FasterTransformers", "djl_python.faster_transformers")
@@ -220,6 +234,8 @@ def prepare_container_def(
             region_name = self.sagemaker_session.boto_session.region_name
             self.image_uri = self.serving_image_uri(region_name)
 
+        environment = self._get_container_env()
+
         local_download_dir = (
             None
             if self.sagemaker_session.settings is None
@@ -258,7 +274,7 @@ def prepare_container_def(
                 kms_key=self.model_kms_key,
             )
             return sagemaker.container_def(
-                self.image_uri, model_data_url=uploaded_code.s3_prefix, env=self.env
+                self.image_uri, model_data_url=uploaded_code.s3_prefix, env=environment
             )
 
     def generate_serving_properties(self, serving_properties={}) -> Dict[str, str]:
@@ -297,6 +313,18 @@ def serving_image_uri(self, region_name):
             version=self.djl_version,
         )
 
+    def _get_container_env(self):
+        if not self.container_log_level:
+            return self.env
+
+        if self.container_log_level not in LOG_LEVEL_MAP:
+            logger.warning(f"Ignoring invalid container log level: {self.container_log_level}")
+            return self.env
+
+        self.env["SERVING_OPTS"] = f'"-Dai.djl.logging.level={LOG_LEVEL_MAP[self.container_log_level]}"'
+        return self.env
+
+
 
 def _determine_engine_for_model_type(model_type: str):
     if model_type in defaults.DEEPSPEED_RECOMMENDED_ARCHITECTURES: