Lint

Signed-off-by: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com>

Author: arekay-nv

src/inference_endpoint/commands/benchmark.py

@@ -59,6 +59,7 @@
 from inference_endpoint.config.schema import (
     Dataset as DatasetConfig,
 )
+from inference_endpoint.dataset_manager.predefined.random import RandomDataset
 from inference_endpoint.config.yaml_loader import ConfigError, ConfigLoader
 from inference_endpoint.core.types import QueryResult
 from inference_endpoint.dataset_manager.dataset import Dataset
@@ -458,7 +459,6 @@ def _run_benchmark(
             raise InputValidationError(
                 "A tokenizer is required to generate the warmup dataset. Ensure model_params.name is set."
             )
-        from inference_endpoint.dataset_manager.predefined.random import RandomDataset
 
         warmup_cfg = config.warmup
         try:

src/inference_endpoint/load_generator/session.py

@@ -339,8 +339,8 @@ def start(
             warmup_generator = load_generator_cls(
                 sample_issuer,
                 warmup_dataset,
-                warmup_sched, # type: ignore[arg-type]
-                *args,  
+                warmup_sched,  # type: ignore[arg-type]
+                *args,
             )
 
         # Create accuracy test generators

tests/integration/commands/test_warmup.py

@@ -33,7 +33,7 @@ def mock_tokenizer():
 
 
 @pytest.mark.integration
-@pytest.mark.asyncio
+@pytest.mark.asyncio(mode="strict")
 async def test_warmup_offline_with_echo_server(
     mock_http_echo_server, ds_pickle_dataset_path, tmp_path, caplog, mock_tokenizer
 ):
@@ -108,7 +108,7 @@ async def test_warmup_offline_with_echo_server(
 
 
 @pytest.mark.integration
-@pytest.mark.asyncio
+@pytest.mark.asyncio(mode="strict")
 async def test_warmup_online_with_echo_server(
     mock_http_echo_server, ds_pickle_dataset_path, tmp_path, caplog, mock_tokenizer
 ):