Repository: LMCache/LMCache
File: benchmarks/multi_round_qa/multi-round-qa.py (Python)

Code before:
            )
        else:
            user_session = UserSession(user_config, self.use_sharegpt)
        self.sessions.append(user_session)
        return user_session

    def _remove_finished_sessions(self):
        sessions_to_remove = [s for s in self.sessions if s.finished]
        if len(sessions_to_remove) > 0:
            logger.info(
                f"Removing {len(sessions_to_remove)} finished sessions, now "
                f"active users: {len(self.sessions) - len(sessions_to_remove)}"
            )
            for session in sessions_to_remove:
                self.session_summaries.append(session.summary())
        self.sessions = [s for s in self.sessions if not s.finished]

    def step(self, timestamp: float, executor: RequestExecutor):
        if self.need_ramp_up:
            self._ramp_up(timestamp, self.ramp_up_time)

        if self.start_time is None:
            self.start_time = timestamp

        # New user session only joins when meets time interval and active user count is less than configured
        if timestamp - self.last_user_join > self.gap_between_users and len(self.sessions) < self.workload_config.num_users:
            self._create_user_session()
            self.last_user_join = timestamp
            logger.info(
                f"Joined a new user {self.user_id}, "
                f"now active users: {len(self.sessions)}"
            )

        for session in self.sessions:
            session.step(timestamp, executor)

        self._remove_finished_sessions()

    @staticmethod
    def ProcessSummary(
        df: pd.DataFrame,
        start_time: Optional[float] = None,
        end_time: Optional[float] = None,
        pending_queries: int = 0,
        config_qps: Optional[float] = None,
    ):
        if start_time and end_time:
            launched_queries = len(
                df.query(f"{start_time} <= launch_time <= {end_time}")
            )
            df = df.query(f"{start_time} <= finish_time <= {end_time}")

Review comment (performance):
In my opinion, we should not add this additional condition. We use this script to benchmark different engines. If an engine performs poorly and requests finish slowly, adding this condition will cause the QPS to decrease, making the workload lighter than for an engine that performs well. This would result in an unfair comparison.

Code after:

    def _create_user_session(self):
        self.user_id += 1
        user_config = UserConfig.new_user_config(self.user_id, self.workload_config)
        if self.use_sharegpt:
            user_session = UserSession(
                user_config, self.use_sharegpt, self.sharegpt_data[self.user_id]
            )
        else:
            user_session = UserSession(user_config, self.use_sharegpt)
        self.sessions.append(user_session)
        return user_session

    def _remove_finished_sessions(self):
        sessions_to_remove = [s for s in self.sessions if s.finished]
        if len(sessions_to_remove) > 0:
            logger.info(
                f"Removing {len(sessions_to_remove)} finished sessions, now "
                f"active users: {len(self.sessions) - len(sessions_to_remove)}"
            )
            for session in sessions_to_remove:
                self.session_summaries.append(session.summary())
        self.sessions = [s for s in self.sessions if not s.finished]

    def _can_join_user(self, timestamp: float) -> bool:
        # No new user session if gap_between_users time interval not meets
        if timestamp - self.last_user_join <= self.gap_between_users:
            return False

        # No user seession if active user count is less than configured
        if (
            self.enforce_strict_concurrent_users
            and len(self.sessions) >= self.workload_config.num_users
        ):
            return False
        return True

    def step(self, timestamp: float, executor: RequestExecutor):
        if self.need_ramp_up:
            self._ramp_up(timestamp, self.ramp_up_time)

        if self.start_time is None:
            self.start_time = timestamp

        # Check if can join new user session
        if self._can_join_user(timestamp):
            self._create_user_session()
            self.last_user_join = timestamp
            logger.info(
                f"Joined a new user {self.user_id}, "
                f"now active users: {len(self.sessions)}"