Use PREFILL_SLICES and DECODE_SLICES from envs.py in disagg_utils

Signed-off-by: Xing Liu <xingliu14@gmail.com>

Author: xingliu14

tpu_inference/core/disagg_utils.py

@@ -3,6 +3,8 @@
 import os
 from typing import Tuple
 
+from tpu_inference import envs
+
 PREFILL_SLICES = 'PREFILL_SLICES'
 DECODE_SLICES = 'DECODE_SLICES'
 
@@ -11,7 +13,7 @@ def is_disagg_enabled() -> bool:
     # We triggrer our code path as long as prefill slices are set. This
     # allows us to test interleave mode effectively with the code path
     # for comparison purposes.
-    return PREFILL_SLICES in os.environ
+    return bool(envs.PREFILL_SLICES)
 
 
 def _parse_slices(slices_str: str) -> Tuple[int, ...]:
@@ -40,12 +42,12 @@ def _parse_slices(slices_str: str) -> Tuple[int, ...]:
 
 
 def get_prefill_slices() -> Tuple[int, ...]:
-    if PREFILL_SLICES not in os.environ:
+    if not envs.PREFILL_SLICES:
         return ()
-    return _parse_slices(os.environ[PREFILL_SLICES])
+    return _parse_slices(envs.PREFILL_SLICES)
 
 
 def get_decode_slices() -> Tuple[int, ...]:
-    if DECODE_SLICES not in os.environ:
+    if not envs.DECODE_SLICES:
         return ()
-    return _parse_slices(os.environ[DECODE_SLICES])
+    return _parse_slices(envs.DECODE_SLICES)