Repository: microsoft/RD-Agent
File: rdagent/utils/workflow.py (Python)

Code before:
                with logger.tag(f"Loop_{li}.{name}"):
                    start = datetime.datetime.now(datetime.timezone.utc)
                    func: Callable[..., Any] = cast(Callable[..., Any], getattr(self, name))
                    try:
                        self.loop_prev_out[name] = func(self.loop_prev_out)
                        # TODO: Fix the error logger.exception(f"Skip loop {li} due to {e}")
                    except Exception as e:
                        if isinstance(e, self.skip_loop_error):
                            # FIXME: This does not support previous demo (due to their last step is not for recording)
                            logger.warning(f"Skip loop {li} due to {e}")
                            # NOTE: strong assumption!  The last step is responsible for recording information
                            self.step_idx = len(self.steps) - 1  # directly jump to the last step.
                            self.loop_prev_out[self.EXCEPTION_KEY] = e
                            continue
                        elif isinstance(e, self.withdraw_loop_error):
                            logger.warning(f"Withdraw loop {li} due to {e}")
                            # Reload the last successful state for this step
                            # Construct the path to the last dump
                            prev_session_dir = self.session_folder / str(li - 1)
                            prev_path = min(
                                (
                                    (int(p.name.split("_", 1)[0]), p)
                                    for p in prev_session_dir.glob("*_*")
                                    if p.is_file()
                                ),
                                key=lambda item: item[0],
                                default=(None, None),
                            )[1]
                            if prev_path:
                                loaded = type(self).load(
                                    prev_path,
                                    output_path=self.session_folder.parent,
                                    do_truncate=False,
                                    replace_timer=True,
                                )
                                # Overwrite current instance state
                                self.__dict__.update(loaded.__dict__)
                                # Continue to next iteration (will restart current loop)
                                continue
                            else:
                                logger.error(f"No previous dump found at {prev_session_dir}, cannot withdraw loop {li}")
                                raise
                        else:
                            raise
                    finally:
                        # make sure failure steps are displayed correclty
                        end = datetime.datetime.now(datetime.timezone.utc)
                        self.loop_trace[li].append(LoopTrace(start, end, step_idx=si))

                        # Update tqdm progress bar directly to step_idx
                        pbar.n = si + 1

Review comment (suggestion):
Why not 
```suggestion
                                key=lambda item: int(p.name.split("_", 1)[0]),
```

and ` p for p in prev_session_dir.glob("*_*") if p.is_file()`

Code after:
                with logger.tag(f"Loop_{li}.{name}"):
                    start = datetime.datetime.now(datetime.timezone.utc)
                    func: Callable[..., Any] = cast(Callable[..., Any], getattr(self, name))
                    try:
                        self.loop_prev_out[name] = func(self.loop_prev_out)
                        # TODO: Fix the error logger.exception(f"Skip loop {li} due to {e}")
                    except Exception as e:
                        if isinstance(e, self.skip_loop_error):
                            # FIXME: This does not support previous demo (due to their last step is not for recording)
                            logger.warning(f"Skip loop {li} due to {e}")
                            # NOTE: strong assumption!  The last step is responsible for recording information
                            self.step_idx = len(self.steps) - 1  # directly jump to the last step.
                            self.loop_prev_out[self.EXCEPTION_KEY] = e
                            continue
                        elif isinstance(e, self.withdraw_loop_error):
                            logger.warning(f"Withdraw loop {li} due to {e}")
                            # Back to previous loop
                            self.step_backward(li - 1)
                            continue
                        else:
                            raise
                    finally:
                        # make sure failure steps are displayed correclty
                        end = datetime.datetime.now(datetime.timezone.utc)
                        self.loop_trace[li].append(LoopTrace(start, end, step_idx=si))

                        # Update tqdm progress bar directly to step_idx
                        pbar.n = si + 1
                        pbar.set_postfix(
                            loop_index=li, step_index=si + 1, step_name=name
                        )  # step_name indicate  last finished step_name

                # index increase and save session
                self.step_idx = (self.step_idx + 1) % len(self.steps)
                if self.step_idx == 0:  # reset to step 0 in next round
                    self.loop_idx += 1
                    if loop_n is not None:
                        loop_n -= 1
                    self.loop_prev_out = {}
                    pbar.reset()  # reset the progress bar for the next loop

                self.dump(self.session_folder / f"{li}" / f"{si}_{name}")  # save a snapshot after the session

    def step_backward(self, li: int) -> None:
        prev_session_dir = self.session_folder / str(li)
        prev_path = min(
            (p for p in prev_session_dir.glob("*_*") if p.is_file()),
            key=lambda item: int(item.name.split("_", 1)[0]),
            default=None,
        )
        if prev_path: