app.py

import datetime
import os
import subprocess
import gradio as gr


CUSTOM_CSS = """
#output_box textarea {
    font-family: IBM Plex Mono, ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, "Liberation Mono", "Courier New", monospace;
}
#run_button {
    visibility: hidden;
}
"""

CUSTOM_JS_ON_LOAD = """
async () => {
    console.log("page.load");
    window.itv = setInterval(() => {
        const root = document.querySelector('body > gradio-app').shadowRoot || document.querySelector('body > gradio-app');
        // ≠ in local or in Spaces for some reason
        const btn = root.querySelector("#run_button");
        btn.click();
    }, 1000);
    // otherwise JS error
    return [];
}
"""


def run():
    output: str = ""
    try:
        output = subprocess.check_output(["nvidia-smi"], text=True)
    except FileNotFoundError:
        output = subprocess.check_output(["ls", "-alh"], text=True)
    comment = (
        datetime.datetime.now().replace(microsecond=0).isoformat().replace("T", " ")
    )
    return f"# {comment}\n\n{output}"


def run_custom_command(custom_command: str, secret: str):
    if secret != os.environ.get("SECRET"):
        return "You can't access this"

    print("custom_command", custom_command)
    try:
        return subprocess.check_output(custom_command.split(), text=True)
    except Exception as e:
        return f"{e}"


output = gr.Textbox(
    label="Command Output", max_lines=32, elem_id="output_box", value=run()
)

with gr.Blocks(css=CUSTOM_CSS) as demo:
    gr.Markdown("#### `nvidia-smi`: How is my GPU Space running right now 🔥")

    with gr.Accordion(label="Power user mode", open=False):
        custom_command = gr.Textbox(label="Input command", value="pwd")
        secret = gr.Textbox(
            label="Secret",
        )
        custom_command_btn = gr.Button("Run")
        custom_command_btn.click(
            fn=run_custom_command,
            inputs=[custom_command, secret],
            outputs=output,
        )
        custom_command_btn.click(
            fn=None,
            inputs=[],
            outputs=[],
            _js="(...args) => { console.log('itv', window.itv, args); clearInterval(window.itv); return args; }",
        )

    output.render()

    run_btn = gr.Button("Run", elem_id="run_button")
    run_btn.click(fn=run, inputs=[], outputs=output)
    demo.load(_js=CUSTOM_JS_ON_LOAD)


demo.launch()