Fix all the lint issues.

Author: wang2yn84

benchmarks/run_offline.py

@@ -109,7 +109,7 @@ def main(argv):
 
   profiling_output = FLAGS.profiling_output
   print("======= decode starting ===")
-  
+
   dec_times = []
   for i in range(10):
     if profiling_output and i == 7 and not profiler_started:

jetstream_pt/layers.py

@@ -456,7 +456,7 @@ def attend(xq, keys, values, local_mask=None):
       # When GQA is enabled, it not necessary to expand
       if not (self.env.ragged_mha and n_rep > 1) and seqlen == 1:
         true_len = 2
-        #xq = torch.broadcast_to(xq, (xq.shape[0], xq.shape[1], 2, xq.shape[3]))
+        # xq = torch.broadcast_to(xq, (xq.shape[0], xq.shape[1], 2, xq.shape[3]))
         xq = torch.nn.functional.pad(
             xq, (0, 0, 0, true_len - seqlen), "constant", 0
         )
@@ -714,6 +714,7 @@ def attend(xq, keys, values, k_scaler, v_scaler, local_mask=None):
 
       return attn_out
 
+
 class Attention(ModuleBase):
   """Attention module."""
 

run_interactive.py

@@ -18,13 +18,10 @@
 from typing import List
 
 # import torch_xla2 first!
-import torch_xla2  # pylint: disable
 import jax
 import numpy as np
-from absl import app, flags
-from colorama import Fore, Style
+from absl import app
 from jetstream.engine import token_utils
-from jetstream_pt import engine as je
 from jetstream_pt.config import FLAGS, create_engine_from_config_flags
 
 
@@ -54,10 +51,15 @@ def main(argv):
   if profiling_prefill:
     jax.profiler.stop_trace()
   prompts: List[str] = [
+      # pylint: disable-next=all
       "I believe the meaning of life is",
+      # pylint: disable-next=all
       "To add an element to an ArrayList of a specific class type in Java, you can follow the following steps:\n\n1. Create an instance of the class to be added.\n2. Get a reference to the ArrayList.\n3. Call the `add()` method on the ArrayList, passing the instance of the class as the argument.\n\nHere's an example of how to add an object of type `Person` to an ArrayList of type `ArrayList<Person>`:\n```csharp\n// Create a new instance of the Person class\nPerson person = new Person(\"John\", 25);\n\n// Get a reference to the ArrayList\nArrayList<Person> peopleList = new ArrayList<>();\n\n// Add the person object to the ArrayList\npeopleList.add(person);\n```\nIn this example, the `Person` class is assumed to have a constructor that takes two arguments: a String for the person's name, and an int for their age. You can substitute your own class and constructor as necessary.",
+      # pylint: disable-next=all
       "<s>[INST] <<SYS>>\nYou are an AI assistant. User will you give you a task. Your goal is to complete the task as faithfully as you can. While performing the task think step-by-step and justify your steps.\n<</SYS>>\n\nQuestion 1: What is commercial real estate finance?\nQuestion 2: What are Commercial Real Estate services?\nOptions are:\n[a]. no.\n[b]. yes.\nWould the answer to these two questions be the same? [/INST]",
+      # pylint: disable-next=all
       "<s>[INST] <<SYS>>\nYou are an AI assistant that helps people find information. Provide a detailed answer so user don\u2019t need to search outside to understand the answer.\n<</SYS>>\n\nUse reasoning to lead to the answer of the following question:\nWhere are you likely to find water underneath?\nOptions:\n- toilet\n- sink\n- jar\n- bridge\n- house\n Reasoning process: [/INST",
+      # pylint: disable-next=all
       "<s>[INST] <<SYS>>\nYou are an AI assistant. You will be given a task. You must generate a detailed and long answer.\n<</SYS>>\n\nContinue the following story.\n\nKay didn't have shoes that fit her feet properly. She only wore sneakers, because the \nChoose from: [I] shoes  fitted badly. [II] sneakers  fitted badly. [/INST]",
   ]
   for prompt in prompts:

run_interactive_disaggregated.py

@@ -19,9 +19,7 @@
 from typing import List
 from absl import app
 from absl import flags
-from colorama import Fore, Style
 
-import numpy as np
 import jax
 
 from jetstream.engine import token_utils
@@ -129,7 +127,6 @@ def main(argv):
   print("Load params ", time.perf_counter() - start)
 
   metadata = prefill_engine.get_tokenizer()
-  tokenizer = prefill_engine.build_tokenizer(metadata)
   vocab = token_utils.load_vocab(metadata.path, metadata.extra_ids)
   stop_tokens = [vocab.eos_id, vocab.pad_id]
   max_output_length = 1024
@@ -157,19 +154,21 @@ def main(argv):
     print(f"---- Input prompts are: {prompt}")
     print(f"---- Encoded tokens are: {tokens}")
 
-    # pylint: disable-next=all
     print(
+        # pylint: disable-next=all
         f"---- Do prefill in prefill engine pod_slice_name: {prefill_engine.pod_slice_name}"
     )
     prefill_result, _ = prefill_engine.prefill(
         params=None, padded_tokens=tokens, true_length=true_length
     )
     print(
+        # pylint: disable-next=all
         f"---- Transfer prefill result to decode engine pod_slice_name: {decode_engine.pod_slice_name}"
     )
     decode_engine.transfer(prefill_result)
-    # pylint: disable-next=all
+
     print(
+        # pylint: disable-next=all
         f"---- Do insert in decode engine pod_slice_name: {decode_engine.pod_slice_name}"
     )
     decode_state = decode_engine.insert(prefill_result, None, slot=slot)

run_interactive_multiple_host.py

@@ -19,7 +19,6 @@
 
 import jax
 from absl import app, flags
-from colorama import Fore, Style
 from jetstream.engine import token_utils
 from jetstream_pt import ray_engine
 from jetstream_pt.config import FLAGS

run_ray_serve_interleave.py

@@ -40,6 +40,7 @@
 
 
 def create_head_resource_name(generation, tpu_chips):
+  """Create head resource name."""
   return f"TPU-{generation}-{tpu_chips}-head"
 
 
@@ -73,6 +74,7 @@ def create_engine(**kwargs):
 
 @serve.deployment
 class JetStreamDeployment:
+  """JetStream deployment."""
 
   def __init__(self, **kwargs):
     os.environ["XLA_FLAGS"] = (
@@ -111,18 +113,24 @@ def __init__(self, **kwargs):
 
     print("Started jetstream driver....")
 
+  # pylint: disable-next=all
   async def Decode(
-      self, request: jetstream_pb2.DecodeRequest
+      self,
+      # pylint: disable-next=all
+      request: jetstream_pb2.DecodeRequest,
+      # pylint: disable-next=all
   ) -> AsyncIterator[jetstream_pb2.DecodeResponse]:
-
+    """Async decode function."""
     return self.orchestrator.Decode(request)
 
 
 def main(_argv):
+  """Main function"""
   resource_name = create_head_resource_name(
       FLAGS.tpu_generation, FLAGS.tpu_chips
   )
   print(f"Using head resource {resource_name}")
+  # pylint: disable-next=all
   deployment = JetStreamDeployment.options(
       ray_actor_options={"resources": {resource_name: 1}}
   ).bind(

run_server.py

@@ -17,7 +17,6 @@
 from typing import Sequence
 
 # import torch_xla2 first!
-import torch_xla2  # pylint: disable
 import jax
 from absl import app, flags
 from jetstream.core import server_lib

run_server_with_ray.py

@@ -19,7 +19,6 @@
 from absl import app, flags
 
 # import torch_xla2 first!
-import torch_xla2  # pylint: disable
 import jax
 from jetstream.core import server_lib
 from jetstream.core.config_lib import ServerConfig

tests/helpers.py

@@ -6,6 +6,7 @@
 from jetstream_pt import environment
 
 
+# pylint: disable-next=all
 def make_env_tiny(bf16_enable=True, env_data_update_fn=lambda _: None):
   torch_dtype = torch.bfloat16 if bf16_enable else torch.float32
   torch.set_default_dtype(torch_dtype)
@@ -33,6 +34,7 @@ def make_env_tiny(bf16_enable=True, env_data_update_fn=lambda _: None):
   return env, config
 
 
+# pylint: disable-next=all
 def make_mixtral_env(bf16_enable=True):
   torch_dtype = torch.bfloat16 if bf16_enable else torch.float32
   torch.set_default_dtype(torch_dtype)
@@ -57,14 +59,16 @@ def make_mixtral_env(bf16_enable=True):
   return env, config
 
 
+# pylint: disable-next=all
 def to_xla_tensor(tree):
   return torch_xla2.default_env().to_xla(tree)
 
 
+# pylint: disable-next=all
 def call_xla_model(model, weights, args):
   with jax.default_device(jax.devices("cpu")[0]):
     xla_weights, xla_inputs = to_xla_tensor((weights, args))
     with torch_xla2.default_env():
       result = torch.func.functional_call(model, xla_weights, xla_inputs)
-    result_torch = torch_xla2.tensor.j2t(result._elem)
+    result_torch = torch_xla2.tensor.j2t(result.jax())
     return result_torch

tests/test_hf_names.py

@@ -4,10 +4,13 @@
 
 
 class TestModuleBase(unittest.TestCase):
+  """Test module base."""
 
   def test_get_hf_names_to_real_name(self):
+    """Test get hugginface names to real name."""
 
     class MyModule(ModuleBase):
+      """My module."""
 
       def __init__(self):
         super().__init__()
@@ -18,6 +21,9 @@ def __init__(self):
         self.param = torch.nn.Parameter(torch.randn(10))
         self.hf_name("param", "model.param")
 
+      def forward(self):
+        """Forward function."""
+
     module = MyModule()
     expected_mapping = {
         "model.my_linear1.weight": "linear1.weight",
@@ -30,20 +36,30 @@ def __init__(self):
     self.assertEqual(module.get_hf_names_to_real_name(), expected_mapping)
 
   def test_get_sharding_annotations(self):
+    """Test get sharding annotations."""
+
     class MyModule(ModuleBase):
+      """MyModule."""
 
       def __init__(self):
         super().__init__()
         self.linear = torch.nn.Linear(10, 20)
         self.embedding = torch.nn.Embedding(100, 50)
         self.inner = InnerModule()
 
+      def forward(self):
+        """Forward function."""
+
     class InnerModule(ModuleBase):
+      """Inner modeule."""
 
       def __init__(self):
         super().__init__()
         self.fc = torch.nn.Linear(50, 100)
 
+      def forward(self):
+        """Forward function."""
+
     module = MyModule()
     module.annotate_sharding("linear.weight", 0)
     module.annotate_sharding("embedding.weight", 1)