WalisonCruz/function-gemma
0
1"""Shared functions for the Trackio UI."""2 3import os4from functools import lru_cache5 6import gradio as gr7import huggingface_hub as hf8 9try:10 import trackio.utils as utils11 from trackio.sqlite_storage import SQLiteStorage12 from trackio.ui.components.colored_checkbox import ColoredCheckboxGroup13 from trackio.ui.helpers.run_selection import RunSelection14except ImportError:15 import utils16 from sqlite_storage import SQLiteStorage17 from ui.components.colored_checkbox import ColoredCheckboxGroup18 from ui.helpers.run_selection import RunSelection19 20CONFIG_COLUMN_MAPPINGS = {21 "_Username": "Username",22 "_Created": "Created",23 "_Group": "Group",24}25CONFIG_COLUMN_MAPPINGS_REVERSE = {v: k for k, v in CONFIG_COLUMN_MAPPINGS.items()}26 27 28HfApi = hf.HfApi()29 30 31def get_project_info() -> str | None:32 dataset_id = os.environ.get("TRACKIO_DATASET_ID")33 space_id = utils.get_space()34 if utils.persistent_storage_enabled():35 return "✨ Persistent Storage is enabled, logs are stored directly in this Space."36 if dataset_id:37 sync_status = utils.get_sync_status(SQLiteStorage.get_scheduler())38 upgrade_message = f"New changes are synced every 5 min <span class='info-container'><input type='checkbox' class='info-checkbox' id='upgrade-info'><label for='upgrade-info' class='info-icon'>ⓘ</label><span class='info-expandable'> To avoid losing data between syncs, <a href='https://huggingface.co/spaces/{space_id}/settings' class='accent-link'>click here</a> to open this Space's settings and add Persistent Storage. Make sure data is synced prior to enabling.</span></span>"39 if sync_status is not None:40 info = f"↻ Backed up {sync_status} min ago to <a href='https://huggingface.co/datasets/{dataset_id}' target='_blank' class='accent-link'>{dataset_id}</a> | {upgrade_message}"41 else:42 info = f"↻ Not backed up yet to <a href='https://huggingface.co/datasets/{dataset_id}' target='_blank' class='accent-link'>{dataset_id}</a> | {upgrade_message}"43 return info44 return None45 46 47def get_projects(request: gr.Request):48 projects = SQLiteStorage.get_projects()49 if project := request.query_params.get("project"):50 interactive = False51 else:52 interactive = True53 if selected_project := request.query_params.get("selected_project"):54 project = selected_project55 else:56 project = projects[0] if projects else None57 58 return gr.Dropdown(59 label="Project",60 choices=projects,61 value=project,62 allow_custom_value=True,63 interactive=interactive,64 info=get_project_info(),65 )66 67 68def update_navbar_value(project_dd, request: gr.Request):69 write_token = None70 if hasattr(request, "query_params") and request.query_params:71 write_token = request.query_params.get("write_token")72 73 metrics_url = f"?selected_project={project_dd}"74 media_url = f"media?selected_project={project_dd}"75 runs_url = f"runs?selected_project={project_dd}"76 files_url = f"files?selected_project={project_dd}"77 78 if write_token:79 metrics_url += f"&write_token={write_token}"80 media_url += f"&write_token={write_token}"81 runs_url += f"&write_token={write_token}"82 files_url += f"&write_token={write_token}"83 84 return gr.Navbar(85 value=[86 ("Metrics", metrics_url),87 ("Media & Tables", media_url),88 ("Runs", runs_url),89 ("Files", files_url),90 ]91 )92 93 94@lru_cache(maxsize=32)95def check_hf_token_has_write_access(hf_token: str | None) -> None:96 """97 Checks to see if the provided hf_token is valid and has write access to the Space98 that Trackio is running in. If the hf_token is valid or if Trackio is not running99 on a Space, this function does nothing. Otherwise, it raises a PermissionError.100 """101 if os.getenv("SYSTEM") == "spaces": # if we are running in Spaces102 # check auth token passed in103 if hf_token is None:104 raise PermissionError(105 "Expected a HF_TOKEN to be provided when logging to a Space"106 )107 who = HfApi.whoami(hf_token)108 owner_name = os.getenv("SPACE_AUTHOR_NAME")109 repo_name = os.getenv("SPACE_REPO_NAME")110 # make sure the token user is either the author of the space,111 # or is a member of an org that is the author.112 orgs = [o["name"] for o in who["orgs"]]113 if owner_name != who["name"] and owner_name not in orgs:114 raise PermissionError(115 "Expected the provided hf_token to be the user owner of the space, or be a member of the org owner of the space"116 )117 # reject fine-grained tokens without specific repo access118 access_token = who["auth"]["accessToken"]119 if access_token["role"] == "fineGrained":120 matched = False121 for item in access_token["fineGrained"]["scoped"]:122 if (123 item["entity"]["type"] == "space"124 and item["entity"]["name"] == f"{owner_name}/{repo_name}"125 and "repo.write" in item["permissions"]126 ):127 matched = True128 break129 if (130 (131 item["entity"]["type"] == "user"132 or item["entity"]["type"] == "org"133 )134 and item["entity"]["name"] == owner_name135 and "repo.write" in item["permissions"]136 ):137 matched = True138 break139 if not matched:140 raise PermissionError(141 "Expected the provided hf_token with fine grained permissions to provide write access to the space"142 )143 # reject read-only tokens144 elif access_token["role"] != "write":145 raise PermissionError(146 "Expected the provided hf_token to provide write permissions"147 )148 149 150@lru_cache(maxsize=32)151def check_oauth_token_has_write_access(oauth_token: str | None) -> None:152 """153 Checks to see if the oauth token provided via Gradio's OAuth is valid and has write access154 to the Space that Trackio is running in. If the oauth token is valid or if Trackio is not running155 on a Space, this function does nothing. Otherwise, it raises a PermissionError.156 """157 if not os.getenv("SYSTEM") == "spaces":158 return159 if oauth_token is None:160 raise PermissionError(161 "Expected an oauth to be provided when logging to a Space"162 )163 who = HfApi.whoami(oauth_token)164 user_name = who["name"]165 owner_name = os.getenv("SPACE_AUTHOR_NAME")166 if user_name == owner_name:167 return168 # check if user is a member of an org that owns the space with write permissions169 for org in who["orgs"]:170 if org["name"] == owner_name and org["roleInOrg"] == "write":171 return172 raise PermissionError(173 "Expected the oauth token to be the user owner of the space, or be a member of the org owner of the space"174 )175 176 177def get_group_by_fields(project: str):178 configs = SQLiteStorage.get_all_run_configs(project) if project else {}179 keys = set()180 for config in configs.values():181 keys.update(config.keys())182 keys.discard("_Created")183 keys = [CONFIG_COLUMN_MAPPINGS.get(key, key) for key in keys]184 choices = [None] + sorted(keys)185 return gr.Dropdown(186 choices=choices,187 value=None,188 interactive=True,189 )190 191 192def group_runs_by_config(193 project: str, config_key: str, filter_text: str | None = None194) -> dict[str, list[str]]:195 if not project or not config_key:196 return {}197 display_key = config_key198 config_key = CONFIG_COLUMN_MAPPINGS_REVERSE.get(config_key, config_key)199 configs = SQLiteStorage.get_all_run_configs(project)200 groups: dict[str, list[str]] = {}201 for run_name, config in configs.items():202 if filter_text and filter_text not in run_name:203 continue204 group_name = config.get(config_key, "None")205 label = f"{display_key}: {group_name}"206 groups.setdefault(label, []).append(run_name)207 for label in groups:208 groups[label].sort()209 sorted_groups = dict(sorted(groups.items(), key=lambda kv: kv[0].lower()))210 return sorted_groups211 212 213def run_checkbox_update(selection: RunSelection, **kwargs) -> gr.CheckboxGroup:214 color_palette = utils.get_color_palette()215 return ColoredCheckboxGroup(216 choices=selection.choices,217 value=selection.selected,218 colors=[219 color_palette[i % len(color_palette)] for i in range(len(selection.choices))220 ],221 label=f"Runs ({len(selection.choices)})",222 **kwargs,223 )224 225 226def handle_run_checkbox_change(227 selected_runs: list[str] | None, selection: RunSelection228) -> RunSelection:229 selection.select(selected_runs or [])230 return selection231 232 233def group_checkbox_update(234 group_runs: list[str], selection: RunSelection235) -> ColoredCheckboxGroup:236 color_palette = utils.get_color_palette()237 choice_indices = {run: i for i, run in enumerate(selection.choices)}238 colors = [239 color_palette[choice_indices.get(run, 0) % len(color_palette)]240 for run in group_runs241 ]242 subset = utils.ordered_subset(group_runs, selection.selected)243 return ColoredCheckboxGroup(244 choices=group_runs,245 value=subset,246 colors=colors,247 label=f"Runs ({len(group_runs)})",248 )249 250 251def handle_group_checkbox_change(252 group_selected: list[str] | None,253 selection: RunSelection,254 group_runs: list[str] | None,255):256 selection.replace_group(group_runs or [], group_selected or [])257 return (258 selection,259 group_checkbox_update(group_runs or [], selection),260 run_checkbox_update(selection),261 )262 263 264def handle_group_toggle(265 select_all: bool,266 selection: RunSelection,267 group_runs: list[str] | None,268):269 target = list(group_runs or []) if select_all else []270 selection.replace_group(group_runs or [], target)271 return (272 selection,273 group_checkbox_update(group_runs or [], selection),274 run_checkbox_update(selection),275 )276