🎨 Fix formatting

Signed-off-by: gkumbhat <kumbhat.gaurav@gmail.com>

Author: gkumbhat

aiu_fms_testing_utils/utils/__init__.py

@@ -67,7 +67,6 @@ def warmup_model(
     dprint(f"PT compile complete, took {pt_compile_model_time:.3f}s")
 
 
-
 def get_env_to_int_list(env_var_name, default):
     """Utility function to convert list of strings passed as given environment variable to
     list of integers
@@ -78,7 +77,7 @@ def get_env_to_int_list(env_var_name, default):
     if isinstance(env_var_string, list):
         return env_var_string
 
-    return [int(v) for v in env_var_string.split(",")  if not isinstance(v, int)]
+    return [int(v) for v in env_var_string.split(",") if not isinstance(v, int)]
 
 
 def ids_for_prompt(prompt, tokenizer):

tests/testing/test_compilation.py

@@ -1,13 +1,11 @@
 """This module contains test related to compilation operation"""
 
 # Standard
-import itertools
 import os
 import pytest
 import time
 
 # Third Party
-from torch import distributed as dist
 import torch
 
 # First Party
@@ -20,7 +18,7 @@
     ids_for_prompt,
     get_env_to_int_list,
     sample_sharegpt_requests,
-    warmup_model
+    warmup_model,
 )
 from aiu_fms_testing_utils.utils.aiu_setup import dprint
 
@@ -37,8 +35,12 @@
 common_model_paths = get_env_to_int_list("COMMON_MODEL_NAME", [GRANITE_3p3_8B_INSTRUCT])
 common_batch_sizes = get_env_to_int_list("FMS_TEST_SHAPES_COMMON_BATCH_SIZES", [1])
 common_seq_lengths = get_env_to_int_list("FMS_TEST_SHAPES_COMMON_SEQ_LENGTHS", [64])
-common_max_new_tokens = get_env_to_int_list("FMS_TEST_SHAPES_COMMON_MAX_NEW_TOKENS", [64])
-common_expected_comp_time = get_env_to_int_list("COMMON_COMPILATION_EXPECTED_TIME", [10]) # In minutes
+common_max_new_tokens = get_env_to_int_list(
+    "FMS_TEST_SHAPES_COMMON_MAX_NEW_TOKENS", [64]
+)
+common_expected_comp_time = get_env_to_int_list(
+    "COMMON_COMPILATION_EXPECTED_TIME", [10]
+)  # In minutes
 
 COMMON_SHAPE_TYPE = "dynamic"
 
@@ -98,7 +100,9 @@ def reset_compiler():
 @pytest.mark.parametrize(
     "model_path,batch_size,seq_length,max_new_tokens,expected_comp_time", COMMON_SHAPES
 )
-def test_compilation_time(model_path, batch_size, seq_length, max_new_tokens, expected_comp_time):
+def test_compilation_time(
+    model_path, batch_size, seq_length, max_new_tokens, expected_comp_time
+):
     """Test to validate time taken for model compilation."""
     torch.manual_seed(42)
     torch.set_default_dtype(torch.float16)
@@ -119,7 +123,7 @@ def test_compilation_time(model_path, batch_size, seq_length, max_new_tokens, ex
     model = get_model(
         architecture="hf_pretrained",
         device_type="cpu",
-        data_type= torch.float16,
+        data_type=torch.float16,
         fused_weights=False,
         **model_path_kwargs,
     )
@@ -137,16 +141,10 @@ def test_compilation_time(model_path, batch_size, seq_length, max_new_tokens, ex
     else:
         COMPILE_DYNAMIC_SHAPE = False
 
-    model.compile(
-        backend="sendnn", options={"sendnn.dynamic": COMPILE_DYNAMIC_SHAPE}
-    )
+    model.compile(backend="sendnn", options={"sendnn.dynamic": COMPILE_DYNAMIC_SHAPE})
     warmup_model(
-        model,
-        input_ids,
-        max_new_tokens,
-        COMPILE_DYNAMIC_SHAPE,
-        **extra_kwargs
+        model, input_ids, max_new_tokens, COMPILE_DYNAMIC_SHAPE, **extra_kwargs
     )
     end_time = time.perf_counter()
 
-    assert (end_time - start_time) < expected_comp_time * 60
+    assert (end_time - start_time) < expected_comp_time * 60