Update serving.py

- Added a new function 'configure_logging()' to dynamically set log levels
- Utilize TS_LOGLEVEL environment variable to control logging verbosity
- Support log levels: OFF, FATAL, ERROR, WARN, INFO, DEBUG, TRACE
- Modify log4j2.xml file using sed command based on TS_LOGLEVEL
- Handle potential errors during log configuration gracefully
- Call configure_logging() before starting TorchServe
- Aim to reduce excessive logging and associated CloudWatch costs
- Maintain default logging if TS_LOGLEVEL is not set or invalid

Author: lavaraja

src/sagemaker_pytorch_serving_container/serving.py

@@ -12,6 +12,8 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+import os
+import subprocess
 from subprocess import CalledProcessError
 
 from retrying import retry
@@ -20,11 +22,35 @@
 
 HANDLER_SERVICE = handler_service.__file__
 
+def configure_logging():
+    log_levels = {
+        '0': 'off',
+        '10': 'fatal',
+        '20': 'error',
+        '30': 'warn',
+        '40': 'info',
+        '50': 'debug',
+        '60': 'trace'
+    }
+
+    ts_loglevel = os.environ.get('TS_LOGLEVEL')
+
+    if ts_loglevel is not None:
+        if ts_loglevel in log_levels:
+            try:
+                log_level = log_levels[ts_loglevel]
+                subprocess.run(['sed', '-i', f's/info/{log_level}/g', 'etc/log4j2.xml'], check=True)
+                print(f"Logging level set to {log_level}")
+            except subprocess.CalledProcessError:
+                print("Error configuring the logging", file=sys.stderr)
+        else:
+            print(f"Invalid TS_LOGLEVEL value: {ts_loglevel}. No changes made to logging configuration.", file=sys.stderr)
+    else:
+        print("TS_LOGLEVEL not set. Using default logging configuration.")
 
 def _retry_if_error(exception):
     return isinstance(exception, CalledProcessError)
 
-
 @retry(stop_max_delay=1000 * 30,
        retry_on_exception=_retry_if_error)
 def _start_torchserve():
@@ -33,6 +59,9 @@ def _start_torchserve():
     # retry starting mms until it's ready
     torchserve.start_torchserve(handler_service=HANDLER_SERVICE)
 
-
 def main():
+    configure_logging()
     _start_torchserve()
+
+if __name__ == '__main__':
+    main()