optimize

app.py

@@ -19,7 +19,7 @@ logging.getLogger("absl").setLevel(logging.WARNING)
 
 API = HfApi(token=os.environ.get("TOKEN"))
 RESULTS_REPO = "open-rl-leaderboard/results"
-REFRESH_RATE = 0.5 * 60  # 5 minutes
+REFRESH_RATE = 5 * 60  # 5 minutes
 ALL_ENV_IDS = {
     "Atari": [
         "AdventureNoFrameskip-v4",
@@ -126,6 +126,7 @@ def iqm(x):
 
 
 def get_leaderboard_df():
+    logger.info("Downloading results")
     dir_path = API.snapshot_download(repo_id=RESULTS_REPO, repo_type="dataset")
     pattern = os.path.join(dir_path, "**", "results_*.json")
     filenames = glob.glob(pattern, recursive=True)
@@ -169,70 +170,43 @@ def format_df(df: pd.DataFrame):
     return df.values.tolist()
 
 
-def refresh_hidden_df():
-    df = get_leaderboard_df()
-    return df
+def refresh_video(df, env_id):
+    env_df = select_env(df, env_id)
+    if not env_df.empty:
+        user_id = env_df.iloc[0]["user_id"]
+        model_id = env_df.iloc[0]["model_id"]
+        model_sha = env_df.iloc[0]["model_sha"]
+        repo_id = f"{user_id}/{model_id}"
+        video_path = API.hf_hub_download(repo_id=repo_id, filename="replay.mp4", revision=model_sha, repo_type="model")
+        return video_path
+    else:
+        return None
 
 
-def refresh_dataframes(df):
-    all_dfs = [format_df(select_env(df, env_id)) for env_id in all_env_ids]
-    return all_dfs
-
-
-def refresh_videos(df):
-    outputs = []
-    for env_id in all_env_ids:
-        env_df = select_env(df, env_id)
-        if not env_df.empty:
-            user_id = env_df.iloc[0]["user_id"]
-            model_id = env_df.iloc[0]["model_id"]
-            model_sha = env_df.iloc[0]["model_sha"]
-            repo_id = f"{user_id}/{model_id}"
-            video_path = API.hf_hub_download(repo_id=repo_id, filename="replay.mp4", revision=model_sha, repo_type="model")
-            outputs.append(video_path)
-        else:
-            outputs.append(None)
-    return outputs
-
-
-def refresh_video(env_id):
-    def func(df):
-        env_df = select_env(df, env_id)
-        if not env_df.empty:
-            user_id = env_df.iloc[0]["user_id"]
-            model_id = env_df.iloc[0]["model_id"]
-            model_sha = env_df.iloc[0]["model_sha"]
-            repo_id = f"{user_id}/{model_id}"
-            video_path = API.hf_hub_download(repo_id=repo_id, filename="replay.mp4", revision=model_sha, repo_type="model")
-            return video_path
-        return None
+def refresh_one_video(df, env_id):
+    def inner():
+        return refresh_video(df, env_id)
 
-    return func
+    return inner
 
 
-def refresh_winners(df):
-    outputs = []
-    for env_id in all_env_ids:
-        env_df = select_env(df, env_id)
-        if not env_df.empty:
-            user_id = env_df.iloc[0]["user_id"]
-            model_id = env_df.iloc[0]["model_id"]
-            url = f"https://huggingface.co/{user_id}/{model_id}"
-            outputs.append(
-                f"""## {env_id}
+def refresh_winner(df, env_id):
+    # print("Refreshing winners")
+    env_df = select_env(df, env_id)
+    if not env_df.empty:
+        user_id = env_df.iloc[0]["user_id"]
+        model_id = env_df.iloc[0]["model_id"]
+        url = f"https://huggingface.co/{user_id}/{model_id}"
+        return f"""## {env_id}
 
 ### 🏆 [Best model]({url}) 🏆"""
-            )
-        else:
-            outputs.append(
-                f"""## {env_id}
+    else:
+        return f"""## {env_id}
 
 This leaderboard is quite empty... 😢
 
 Be the first to [submit your model]()!
 """
-            )
-    return outputs
 
 
 HEADING = """
@@ -332,61 +306,73 @@ h3 {
 
 """
 
+
+def update_globals():
+    global dataframes, winner_texts, video_pathes, df
+    df = get_leaderboard_df()
+    all_env_ids = [env_id for env_ids in ALL_ENV_IDS.values() for env_id in env_ids]
+    dataframes = {env_id: format_df(select_env(df, env_id)) for env_id in all_env_ids}
+    winner_texts = {env_id: refresh_winner(df, env_id) for env_id in all_env_ids}
+    video_pathes = {env_id: refresh_video(df, env_id) for env_id in all_env_ids}
+
+
+update_globals()
+
+
+def refresh():
+    global dataframes, winner_texts, video_pathes
+    return list(dataframes.values()) + list(winner_texts.values()) + [list(video_pathes.values())[0]]
+
+
 with gr.Blocks(css=css) as demo:
     gr.Markdown(HEADING)
     with gr.Tabs(elem_classes="tab-buttons") as tabs:
         with gr.TabItem("🏅 Leaderboard"):
-            df = get_leaderboard_df()
-            hidden_df = gr.components.Dataframe(df, visible=False)
-            all_env_ids = []
-            all_gr_dfs = []
-            all_gr_videos = []
-            all_gr_winners = []
+            all_gr_dfs = {}
+            all_gr_winners = {}
+            all_gr_videos = {}
             for env_domain, env_ids in ALL_ENV_IDS.items():
                 with gr.TabItem(env_domain):
                     for env_id in env_ids:
-                        # If the env_id envs with "NoFrameskip-v4", we remove it
+                        # If the env_id envs with "NoFrameskip-v4", we remove it to improve readability
                         tab_env_id = env_id[: -len("NoFrameskip-v4")] if env_id.endswith("NoFrameskip-v4") else env_id
                         with gr.TabItem(tab_env_id) as tab:
                             logger.info(f"Creating tab for {env_id}")
                             with gr.Row(equal_height=False):
                                 with gr.Column(scale=3):
-                                    # Display the leaderboard
                                     gr_df = gr.components.Dataframe(
                                         headers=["🏆", "🧑 User", "🤖 Model id", "📊 IQM episodic return"],
                                         datatype=["number", "markdown", "markdown", "number"],
-                                        row_count=(20, "fixed"),
                                     )
                                 with gr.Column(scale=1):
                                     with gr.Row():  # Display the env_id and the winner
                                         gr_winner = gr.Markdown()
                                     with gr.Row():  # Play the video of the best model
-                                        gr_video = gr.PlayableVideo(  # Doesn't loop for the moment, see https://github.com/gradio-app/gradio/issues/7689
+                                        gr_video = gr.PlayableVideo(  # Doesn't loop for the moment, see https://github.com/gradio-app/gradio/issues/7689,
                                             min_width=50,
-                                            autoplay=True,
                                             show_download_button=False,
                                             show_share_button=False,
                                             show_label=False,
+                                            interactive=False,
                                         )
 
-                                all_env_ids.append(env_id)
-                                all_gr_dfs.append(gr_df)
-                                all_gr_winners.append(gr_winner)
-                                all_gr_videos.append(gr_video)
+                            all_gr_dfs[env_id] = gr_df
+                            all_gr_winners[env_id] = gr_winner
+                            all_gr_videos[env_id] = gr_video
+
+                            tab.select(refresh_one_video(df, env_id), outputs=[gr_video])
 
-                        tab.select(refresh_video(env_id), inputs=hidden_df, outputs=gr_video)
+                # Load the first video of the first environment
+                demo.load(refresh_one_video(df, env_ids[0]), outputs=[all_gr_videos[env_ids[0]]])
 
         with gr.TabItem("📝 About"):
             gr.Markdown(ABOUT_TEXT)
 
-    demo.load(refresh_hidden_df, outputs=hidden_df, every=REFRESH_RATE)
-    demo.load(refresh_dataframes, inputs=hidden_df, outputs=all_gr_dfs, every=REFRESH_RATE)
-    demo.load(refresh_videos, inputs=hidden_df, outputs=all_gr_videos, every=REFRESH_RATE)
-    demo.load(refresh_winners, inputs=hidden_df, outputs=all_gr_winners, every=REFRESH_RATE)
-
+    demo.load(refresh, outputs=list(all_gr_dfs.values()) + list(all_gr_winners.values()))
 
 scheduler = BackgroundScheduler()
 scheduler.add_job(func=backend_routine, trigger="interval", seconds=REFRESH_RATE, max_instances=1)
+scheduler.add_job(func=update_globals, trigger="interval", seconds=REFRESH_RATE, max_instances=1)
 scheduler.start()