chore: black and isort format on codebase (#298)

Formatted the codebase using black and isort.

Author: raresgaia123

infscale/agent/agent.py

@@ -153,7 +153,7 @@ async def update_worker_status(self, message: WorkerStatusMessage) -> None:
 
     def _handle_worker(self, status_msg: WorkerStatusMessage) -> None:
         """Handle status message received from worker.
-        
+
         Whenever DONE, FAILED or TERMINATED is received, worker manager
         will remove that worker from its state.
         When no workers are available for the given job id, cleanup is done
@@ -334,18 +334,17 @@ def _handle_command(self, action: pb2.Action) -> None:
                         MessageType.FINISH_JOB, WorkerStatus.DONE, action.job_id
                     )
                     self.worker_mgr.send(w, msg)
-                    
+
             case CommandAction.CHECK_LOOP:
                 workers = self.worker_mgr.get_workers_by_job_id(action.job_id)
                 failed_wids_str = action.manifest.decode("utf-8")
                 failed_wids_set = ast.literal_eval(failed_wids_str)
-                
+
                 for w in workers.values():
                     msg = Message(
                         MessageType.CHECK_LOOP, failed_wids_set, action.job_id
                     )
                     self.worker_mgr.send(w, msg)
-                    
 
     async def heart_beat(self):
         """Send a heart beat message periodically."""
@@ -411,7 +410,9 @@ def _terminate_workers(self, config: JobConfig) -> None:
         job_id, force_terminate = config.job_id, config.force_terminate
         stop_wrkrs = self.job_mgr.get_workers(job_id, CommandAction.STOP)
 
-        msg_type = MessageType.FORCE_TERMINATE if force_terminate else MessageType.TERMINATE
+        msg_type = (
+            MessageType.FORCE_TERMINATE if force_terminate else MessageType.TERMINATE
+        )
         self.worker_mgr._signal_terminate_wrkrs(job_id, True, stop_wrkrs, msg_type)
 
     async def report(self):

infscale/agent/worker_manager.py

@@ -75,7 +75,7 @@ def add(
         self._workers[worker.pipe.fileno()] = worker
 
         return worker
-                
+
     def has_workers_for_job(self, job_id: str) -> bool:
         """Return True if there are any workers assigned to the given job ID."""
         return any(worker.job_id == job_id for worker in self._workers.values())
@@ -136,7 +136,7 @@ def on_read_ready(
                 # When ConnectionResetError is raised, the pipe is already closed due to worker failure
                 # so we only need to ignore this error.
                 pass
-            
+
     def remove_worker(self, wrk_id: str) -> None:
         """Remove worker related data."""
         for k, v in list(self._workers.items()):

infscale/configs/job.py

@@ -170,13 +170,16 @@ def get_worlds_to_configure(
         updated_worlds = {
             k
             for k in common_keys
-            if (curr_worlds[k].addr != new_worlds[k].addr
+            if (
+                curr_worlds[k].addr != new_worlds[k].addr
                 or curr_worlds[k].data_port != new_worlds[k].data_port
-                or curr_worlds[k].ctrl_port != new_worlds[k].ctrl_port)
+                or curr_worlds[k].ctrl_port != new_worlds[k].ctrl_port
+            )
         }
 
         return deploy_worlds | updated_worlds
-    
+
+
 class ServeConfigHelper:
     """Class for defining helper methods for serve config."""
 
@@ -399,7 +402,7 @@ def is_identical(x: JobConfig, y: JobConfig) -> bool:
     def world_name(world_id: int) -> str:
         """Return world name given a world id."""
         return f"w{world_id}"
-    
+
     @staticmethod
     def get_pipeline_identifiers(new_cfg: JobConfig) -> set[str]:
         """Get pipeline identifiers based on server id."""
@@ -450,30 +453,38 @@ def categorize_workers(
         # select workers that will be affected by workers to be started
         for w, world_info_list in new_config.flow_graph.items():
             for new_world_info in world_info_list:
-                curr_world_info = helper.find_matching_world_info(curr_config, w, new_world_info)
-                helper.pick_workers(update_wrkrs, start_wrkrs, w, new_world_info, curr_world_info)
+                curr_world_info = helper.find_matching_world_info(
+                    curr_config, w, new_world_info
+                )
+                helper.pick_workers(
+                    update_wrkrs, start_wrkrs, w, new_world_info, curr_world_info
+                )
 
         if curr_config is None:
             return start_wrkrs, update_wrkrs, stop_wrkrs
 
         # select workers that will be affected by workers to be stopped
         for w, world_info_list in curr_config.flow_graph.items():
             for new_world_info in world_info_list:
-                curr_world_info = helper.find_matching_world_info(curr_config, w, new_world_info)
-                helper.pick_workers(update_wrkrs, stop_wrkrs, w, new_world_info, curr_world_info)
+                curr_world_info = helper.find_matching_world_info(
+                    curr_config, w, new_world_info
+                )
+                helper.pick_workers(
+                    update_wrkrs, stop_wrkrs, w, new_world_info, curr_world_info
+                )
 
         # due to pervious state, recover workers are included in update workers
         # therefore, recover workers need to be removed from the updated ones.
         update_wrkrs -= recover_wrkrs
 
         return start_wrkrs, update_wrkrs, stop_wrkrs
-    
+
     @staticmethod
     def get_workers_diff(a: JobConfig, b: JobConfig) -> set[str]:
         """Return a set of worker ids diffs based on old and new cfg."""
         old_workers = {worker.id for worker in a.workers}
         new_workers = {worker.id for worker in b.workers}
-        
+
         return old_workers - new_workers
 
     @staticmethod
@@ -509,10 +520,10 @@ def remove_pipeline(config: JobConfig, workers_to_remove: set[str]) -> JobConfig
 
 class JobConfigHelper:
     """Class for defining helper methods for job config."""
+
     def get_server_id(self, config: JobConfig) -> str:
         return next((w.id for w in config.workers if w.is_server), "")
-    
-    
+
     def find_pipeline_nodes(
         self,
         flow_graph: dict[str, list[WorldInfo]],
@@ -577,9 +588,7 @@ def find_pipeline_nodes(
 
         # everything else (except server) is removed
         to_remove = {
-            wid
-            for wid in flow_graph
-            if wid != server_id and wid not in survivors
+            wid for wid in flow_graph if wid != server_id and wid not in survivors
         }
 
         return to_remove
@@ -600,11 +609,13 @@ def pick_workers(
 
         The needles are workers to start or stop and the haystack is
         name and peers.
-        
+
         Also includes peers of `name` if its connection details
         (`addr`, `ctrl_port`, `data_port`) differ from the previous config.
         """
-        if curr_world_info and self.has_connection_changed(curr_world_info, new_world_info):
+        if curr_world_info and self.has_connection_changed(
+            curr_world_info, new_world_info
+        ):
             for peer in new_world_info.peers:
                 res_set.add(peer)
 
@@ -634,15 +645,15 @@ def pick_workers(
                 # because name is already affected by one peer
                 # so we come out of the for-loop
                 break
-                
+
     def has_connection_changed(self, old: WorldInfo, new: WorldInfo) -> bool:
         """Check if worker connection details are changed."""
         return (
-            old.addr != new.addr or
-            old.ctrl_port != new.ctrl_port or
-            old.data_port != new.data_port
+            old.addr != new.addr
+            or old.ctrl_port != new.ctrl_port
+            or old.data_port != new.data_port
         )
-            
+
     def find_matching_world_info(
         self, curr_config: JobConfig | None, w: str, new_world_info: WorldInfo
     ) -> WorldInfo | None:

infscale/controller/controller.py

@@ -243,12 +243,12 @@ async def send_command_to_agent(
         """Send command to agent."""
         agent_context = self.agent_contexts[agent_id]
         context = agent_context.get_grpc_ctx()
-        
+
         kwargs = {
             "type": action.action,
             "job_id": job_id,
         }
-        
+
         # this is needed because checking against an empty set will be False
         if hasattr(action, "failed_wids"):
             kwargs["manifest"] = str(action.failed_wids).encode("utf-8")
@@ -281,7 +281,9 @@ async def job_setup(
         try:
             self.ctrl.handle_job_ports(req)
         except AttributeError:
-            logger.error(f"Failed to process ports. Job: {req.job_id} Agent: {req.agent_id}")
+            logger.error(
+                f"Failed to process ports. Job: {req.job_id} Agent: {req.agent_id}"
+            )
 
         return empty_pb2.Empty()
 
@@ -307,7 +309,9 @@ async def update_wrk_status(
         try:
             await self.ctrl.handle_wrk_status(request)
         except Exception as e:
-            logger.error(f"Failed to process worker status. Worker: {request.worker_id} Status: {request.status} {e}")
+            logger.error(
+                f"Failed to process worker status. Worker: {request.worker_id} Status: {request.status} {e}"
+            )
 
         return empty_pb2.Empty()
 

infscale/controller/ctrl_dtype.py

@@ -39,7 +39,9 @@ class CommandAction(str, Enum):
     UPDATE = "update"  # CLI - Controller update command
     SETUP = "setup"  # ctrl<->agent setup job, assign port numbers to workers
     FINISH_JOB = "finish_job"  # ctrl<->agent action to notify job's completion
-    CHECK_LOOP = "check_loop"  # ctrl<->agent action for workers to check their pipeline loops
+    CHECK_LOOP = (
+        "check_loop"  # ctrl<->agent action for workers to check their pipeline loops
+    )
 
 
 class CommandActionModel(BaseModel):

infscale/controller/job_checker.py

@@ -17,6 +17,7 @@
 """job_checker.py."""
 
 from collections import defaultdict
+
 from infscale import get_logger
 from infscale.common.job_msg import WorkerStatus
 from infscale.configs.job import JobConfig