lint

Author: tdene

megatron/core/inference/engines/dynamic_engine.py

@@ -1318,9 +1318,9 @@ def schedule_requests(self) -> int:
             header = Headers(data[0])
 
             if self.received_stop:
-                assert header == Headers.STOP_ACK, (
-                    "Engine is shutting down. No other messages allowed except STOP_ACK."
-                )
+                assert (
+                    header == Headers.STOP_ACK
+                ), "Engine is shutting down. No other messages allowed except STOP_ACK."
 
             if header == Headers.SUBMIT_REQUEST:
                 request_id, prompt, sampling_params = data[1:]

megatron/core/inference/headers.py

@@ -20,6 +20,7 @@ class Headers(Enum):
     STOP = auto()
     STOP_ACK = auto()
 
+
 class UnknownHeaderError(Exception):
     """A signal with an unrecognized header was received by the coordinator."""
 

megatron/core/inference/inference_client.py

@@ -139,9 +139,7 @@ async def _recv_task(self):
                     )
                     completion_future = self.completion_futures.pop(request_id)
                     if completion_future.done():
-                        logging.warning(
-                            f"Client: The future for {request_id} has been cancelled!"
-                        )
+                        logging.warning(f"Client: The future for {request_id} has been cancelled!")
                         continue
                     completion_future.set_result(DynamicInferenceRequestRecord.deserialize(reply))
                 elif header == Headers.PAUSE_ACK: