Spaces:
Sleeping
Sleeping
File size: 5,212 Bytes
8528842 2a74e89 8528842 bab875d 8528842 bab875d 8528842 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
import json
import gradio as gr
from mcp_mlflow_tools import (
set_tracking_uri,
get_system_info,
list_experiments,
create_experiment,
register_model,
search_runs,
list_registered_models,
get_model_info
)
def create_interface():
with gr.Blocks(title="MLflow MCP Service") as app:
gr.Markdown("# MLflow MCP Service")
gr.Markdown("A service that exposes MLflow functionality through a web interface and API endpoints.")
with gr.Tab("Tracking & System Info"):
with gr.Group():
gr.Markdown("## Set Tracking URI")
uri_input = gr.Textbox(label="MLflow Tracking URI")
uri_output = gr.JSON(label="Result")
uri_button = gr.Button("Set URI")
uri_button.click(
fn=set_tracking_uri,
inputs=uri_input,
outputs=uri_output
)
with gr.Group():
gr.Markdown("## Get System Info")
sys_info_output = gr.JSON(label="System Information")
sys_info_button = gr.Button("Get Info")
sys_info_button.click(
fn=get_system_info,
inputs=[],
outputs=sys_info_output
)
with gr.Tab("Experiment Management"):
with gr.Group():
gr.Markdown("## List Experiments")
exp_list_output = gr.JSON(label="Experiments")
exp_list_button = gr.Button("List Experiments")
exp_list_button.click(
fn=list_experiments,
inputs=[],
outputs=exp_list_output
)
with gr.Group():
gr.Markdown("## Create Experiment")
exp_name_input = gr.Textbox(label="Experiment Name")
exp_tags_input = gr.Textbox(label="Tags (JSON format)", placeholder='{"key": "value"}')
exp_create_output = gr.JSON(label="Result")
exp_create_button = gr.Button("Create Experiment")
def create_exp_with_tags(name, tags_str):
"""Create a new experiment. Given the name and tags"""
try:
tags = json.loads(tags_str) if tags_str else None
except json.JSONDecodeError:
return {"error": True, "message": "Invalid JSON format for tags"}
return create_experiment(name, tags)
exp_create_button.click(
fn=create_exp_with_tags,
inputs=[exp_name_input, exp_tags_input],
outputs=exp_create_output
)
with gr.Tab("Model Registry"):
with gr.Group():
gr.Markdown("## Register Model")
reg_run_id = gr.Textbox(label="Run ID")
reg_artifact_path = gr.Textbox(label="Artifact Path")
reg_model_name = gr.Textbox(label="Model Name")
reg_output = gr.JSON(label="Result")
reg_button = gr.Button("Register Model")
reg_button.click(
fn=register_model,
inputs=[reg_run_id, reg_artifact_path, reg_model_name],
outputs=reg_output
)
with gr.Group():
gr.Markdown("## List Registered Models")
list_models_output = gr.JSON(label="Models")
list_models_button = gr.Button("List Models")
list_models_button.click(
fn=list_registered_models,
inputs=[],
outputs=list_models_output
)
with gr.Group():
gr.Markdown("## Get Model Info")
model_info_name = gr.Textbox(label="Model Name")
model_info_output = gr.JSON(label="Model Information")
model_info_button = gr.Button("Get Info")
model_info_button.click(
fn=get_model_info,
inputs=model_info_name,
outputs=model_info_output
)
with gr.Tab("Run Search"):
with gr.Group():
gr.Markdown("## Search Runs")
search_exp_id = gr.Textbox(label="Experiment ID")
search_filter = gr.Textbox(label="Filter String")
search_order_by = gr.Textbox(label="Order By")
search_max_results = gr.Number(label="Max Results", value=100, precision=0)
search_output = gr.JSON(label="Search Results")
search_button = gr.Button("Search")
search_button.click(
fn=search_runs,
inputs=[search_exp_id, search_filter, search_order_by, search_max_results],
outputs=search_output
)
return app
if __name__ == "__main__":
app = create_interface()
app.launch(mcp_server=True) |