final touches

Dockerfile

@@ -25,3 +25,5 @@ RUN pip install -r requirements.txt
 
 # Run app
 ENTRYPOINT python start_app.py
+
+# PS: Run with a variable SYSTEM=spaces to enable OAuth on Hugging Face Spaces

run_job.py

@@ -32,7 +32,6 @@ DATA_CARD = "# Dataset Card for {dst}\n\nDataset prepared from {src} using\n\n``
 def sql(src: str, dst: str, query: str, config: str = "default", split: str = "train", private: bool = False, dry_run: bool = False):
     import os
     import duckdb
-    from contextlib import nullcontext
     from huggingface_hub import CommitScheduler, DatasetCard
 
     class CommitAndCleanScheduler(CommitScheduler):
@@ -47,25 +46,24 @@ def sql(src: str, dst: str, query: str, config: str = "default", split: str = "t
             for path in self.last_uploaded:
                 path.unlink(missing_ok=True)
 
-    with nullcontext() if dry_run else CommitAndCleanScheduler(repo_id=dst, repo_type="dataset", folder_path="dst", path_in_repo="data", every=0.1, private=private):
-        con = duckdb.connect(":memory:", config=CONFIG)
-        src_kwargs = con.sql(CMD_SRC_KWARGS.format(src=src, config=config, split=split)).df().to_dict(orient="records")
-        if not src_kwargs:
-            raise ValueError(f'Invalid --config "{config}" for dataset "{src}", please select a valid dataset config/subset.')
+    con = duckdb.connect(":memory:", config=CONFIG)
+    src_kwargs = con.sql(CMD_SRC_KWARGS.format(src=src, config=config, split=split)).df().to_dict(orient="records")
+    if not src_kwargs:
+        raise ValueError(f'Invalid --config "{config}" for dataset "{src}", please select a valid dataset config/subset.')
 
-        con.sql((CMD_SRC_DRY_RUN if dry_run else CMD_SRC).format(**src_kwargs[0]))
-        if dry_run:
-            print(f"Sample data from '{src}' that would be written to dataset '{dst}':\n")
-        else:
-            con.sql("PRAGMA enable_progress_bar;")
+    con.sql((CMD_SRC_DRY_RUN if dry_run else CMD_SRC).format(**src_kwargs[0]))
 
-        result = con.sql((CMD_DST_DRY_RUN if dry_run else CMD_DST).format(query=query.rstrip("\n ;")))
-        DatasetCard(DATA_CARD.format(src=src, dst=dst, query=query)).save("dst/README.md")
-        if dry_run:
-            print(result.df().to_markdown())
-        else:
-            print("done")
+    if dry_run:
+        print(f"Sample data from '{src}' that would be written to dataset '{dst}':\n")
+        result = con.sql(CMD_DST_DRY_RUN.format(query=query.rstrip("\n ;")))
+        print(result.df().to_markdown())
+        return
 
+    with CommitAndCleanScheduler(repo_id=dst, repo_type="dataset", folder_path="dst", path_in_repo="data", every=0.1, private=private):
+        con.sql("PRAGMA enable_progress_bar;")
+        result = con.sql(CMD_DST.format(query=query.rstrip("\n ;")))
+        DatasetCard(DATA_CARD.format(src=src, dst=dst, query=query)).save("dst/README.md")
+        print("done")
 
 if __name__ == '__main__':
     fire.Fire(sql)

start_app.py

@@ -1,6 +1,8 @@
+import json
 import os
 import re
 import subprocess
+import time
 import yaml
 
 import gradio as gr
@@ -24,15 +26,23 @@ except Exception:
 
 DRY_RUN = bool(HELP) and bool(m :=re.search("--dry(-|_)run", HELP)) and m.group(0)
 
-def parse_log(line: str, pbars: dict[str, float]):
-    if (percent_match := re.search("\\d+(?:\\.\\d+)?%", line)) and any(c in line.split("%")[1][:10] for c in "|█▌"):
-        [pbars.pop(desc) for desc, percent in pbars.items() if percent == 1.]
-        percent = float(percent_match.group(0)[:-1]) / 100
-        desc = line[:percent_match.start()].strip() or "Progress"
-        pbars[desc] = percent
-        yield ""
-    else:
-        yield line
+def parse_log(line: str, pbars: dict[str, float] = None):
+    if line.startswith("data: {"):
+        data = json.loads(line[len("data: "):])
+        data, timestamp = data["data"], data["timestamp"]
+        if pbars is not None and data.startswith("===== Job started at"):
+            pbars.pop("Starting ⚙️", None)
+            pbars["Running 🏃"] = 0.0
+            return f"[{timestamp}] {data}\n\n"
+        elif pbars is not None and (percent_match := re.search("\\d+(?:\\.\\d+)?%", data)) and any(c in data.split("%")[1][:10] for c in "|█▌"):
+            pbars.pop("Running 🏃", None)
+            [pbars.pop(desc) for desc, percent in pbars.items() if percent == 1.]
+            percent = float(percent_match.group(0)[:-1]) / 100
+            desc = data[:percent_match.start()].strip() or "Progress"
+            pbars[desc] = percent
+        else:
+            return f"[{timestamp}] {data}\n\n"
+    return ""
 
 def dry_run(src, config, split, dst, query):
     if not all([src, config, split, dst, query]):
@@ -77,15 +87,26 @@ def run(src, config, split, dst, query, oauth_token: gr.OAuthToken | None, profi
         pbars = {"Finished with an error ❌": 1.0}
     else:
         job_id = resp.json()["metadata"]["job_id"]
+        pbars = {"Starting ⚙️": 0.0}
+        yield {output_markdown: logs, progress_labels: gr.Label(pbars, visible=bool(pbars))}
         resp = requests.get(
             f"https://huggingface.co/api/jobs/{username}/{job_id}/logs-stream",
+            headers={"Authorization": f"Bearer {token}"},
+            stream=True
         )
-        for line in iter(resp.raw.readline, b""):
-            logs += parse_log(line.decode(), pbars=pbars)
+        for line in resp.iter_lines():
+            logs += parse_log(line.decode("utf-8"), pbars=pbars)
             yield {output_markdown: logs, progress_labels: gr.Label(pbars, visible=bool(pbars))}
-        job_status = requests.get(
-            f"https://huggingface.co/api/jobs/{username}/{job_id}",
-        ).json()
+        job_status = {"status": {"stage": "RUNNING"}}
+        while True:
+            job_status = requests.get(
+                f"https://huggingface.co/api/jobs/{username}/{job_id}",
+                headers={"Authorization": f"Bearer {token}"}
+            ).json()
+            if job_status["status"]["stage"] == "RUNNING":
+                time.sleep(1)
+            else:
+                break
         if job_status["status"]["stage"] == "COMPLETED":
             pbars = {"Finished ✅": 1.0}
         else:
@@ -119,7 +140,8 @@ with gr.Blocks() as demo:
         if DRY_RUN:
             dry_run_button = gr.Button("Dry-Run")
     progress_labels= gr.Label(visible=False, label="Progress")
-    output_markdown = gr.Markdown(label="Output logs")
+    with gr.Accordion("Details", open=False):
+        output_markdown = gr.Markdown(label="Output logs")
     run_button.click(run, inputs=[dataset_dropdown, subset_dropdown, split_dropdown, dst_dropdown, query_textarea], outputs=[progress_labels, output_markdown])
     if DRY_RUN:
         dry_run_button.click(dry_run, inputs=[dataset_dropdown, subset_dropdown, split_dropdown, dst_dropdown, query_textarea], outputs=[progress_labels, output_markdown])