Update offline_lora_inference.py to use tpu_inference.envs

Author: xingliu14

examples/offline_lora_inference.py

@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import os
 import time
 
-import vllm.envs as envs
+import tpu_inference.envs as envs
+import vllm.envs as vllm_envs
 from vllm import LLM, EngineArgs
 from vllm.lora.request import LoRARequest
 from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -55,13 +55,13 @@ def main(args: dict):
         "lora_adapter_3", 3,
         "Username6568/Qwen2.5-3B-Instruct-1_plus_1_equals_3_adapter")
 
-    if envs.VLLM_TORCH_PROFILER_DIR is not None:
+    if vllm_envs.VLLM_TORCH_PROFILER_DIR is not None:
         llm.start_profile()
     start = time.perf_counter()
     outputs = llm.generate(prompt,
                            sampling_params=sampling_params,
                            lora_request=lora_request)
-    if envs.VLLM_TORCH_PROFILER_DIR is not None:
+    if vllm_envs.VLLM_TORCH_PROFILER_DIR is not None:
         llm.stop_profile()
 
     # Print the outputs.
@@ -77,7 +77,7 @@ def main(args: dict):
 
 if __name__ == "__main__":
     # Skip long warmup for local simple test.
-    os.environ['SKIP_JAX_PRECOMPILE'] = '1'
+    envs.environment_variables['SKIP_JAX_PRECOMPILE'] = lambda: True
 
     parser = create_parser()
     args: dict = vars(parser.parse_args())