""" Submit Task tab UI component. """ import json import streamlit as st from src.clients.tidbcloud_client import ComponentType, get_available_components from src.database import db from src.ui.utils import ( cached_get_nodespecs, cached_get_regions, load_workload_template, ) from src.services.workload_templates import ( TPCC_1K_TEMPLATE, TPCC_X_1K_TEMPLATE, build_premium_workload_config, canonicalize_workload_template_or_default, get_workload_templates_for_plan, ) DEFAULT_ORG_ID_PLACEHOLDER = "1369847559692509562" DEFAULT_PROJECT_ID_PLACEHOLDER = "1369847559694040594" @st.fragment def render_tab1(): """Submit Test tab - wrapped in fragment for isolation.""" st.header("Submit New Test Task") # Identity user_email = st.text_input("User Email", placeholder="your.email@pingcap.com") cluster_version = st.text_input("Cluster Version", placeholder="v8.5.0") # Optional org/project overrides (must be above Step 1 per UX request) st.subheader("TiDB Cloud IDs (Optional)") st.caption( "Leave blank to keep defaults. If you set new values, both org_id and " "project_id are required." ) col_org, col_project = st.columns(2) with col_org: org_id_input = st.text_input( "org_id", value="", placeholder=DEFAULT_ORG_ID_PLACEHOLDER, help="TiDB Cloud organization ID (optional).", key="override_org_id", ).strip() with col_project: project_id_input = st.text_input( "project_id", value="", placeholder=DEFAULT_PROJECT_ID_PLACEHOLDER, help="TiDB Cloud project ID (optional).", key="override_project_id", ).strip() # Step 1: Cloud Provider Selection st.subheader("Step 1: Select Cloud Provider") cloud_options = { "AWS": {"enabled": True}, "GCP": {"enabled": True}, "AZURE": {"enabled": True}, "ALIYUN": {"enabled": True}, } env_options = { "dev": {"enabled": False}, "staging": {"enabled": True}, "prod": {"enabled": False}, } available_clouds = [ cloud for cloud, config in cloud_options.items() if config["enabled"] ] disabled_clouds = [ cloud for cloud, config in cloud_options.items() if not config["enabled"] ] available_envs = [env for env, config in env_options.items() if config["enabled"]] disabled_envs = [ env for env, config in env_options.items() if not config["enabled"] ] env = st.selectbox( "Environment", available_envs, help=( f"Currently unavailable: {', '.join(disabled_envs)}" if disabled_envs else None ), ) cloud_provider = st.selectbox( "Cloud Provider", available_clouds, help=( f"Currently unavailable: {', '.join(disabled_clouds)}" if disabled_clouds else None ), ) # Step 2: Region Selection st.subheader("Step 2: Select Region") if cloud_provider: regions = cached_get_regions(cloud_provider) if regions: region_options = { r[ "regionId" ]: f"{r['regionId']} ({r.get('displayName', r['regionId'])})" for r in regions } selected_region_id = st.selectbox( "Region", options=list(region_options.keys()), format_func=lambda x: region_options[x], ) else: st.warning(f"No regions available for {cloud_provider}") selected_region_id = None else: selected_region_id = None # Plan Type Selection st.subheader("Plan Type") plan_type = st.selectbox( "Plan Type", ["Dedicated", "Premium"], help="Dedicated: full component-level config. Premium: Capacity (RCU) only.", ) # Step 3: TiDB NodeSpec Selection (Dedicated only) if plan_type == "Dedicated": st.subheader("Step 3: Select TiDB Specification") tidb_nodespec = None tidb_vcpu = None if plan_type == "Dedicated" and selected_region_id: tidb_specs = cached_get_nodespecs(selected_region_id, ComponentType.TIDB) if tidb_specs: tidb_spec_options = { s[ "nodeSpecKey" ]: f"{s['nodeSpecKey']} ({s['vCpu']}C{s['memorySizeGi']}G)" for s in tidb_specs } # Use session state to persist selection if "tidb_nodespec" not in st.session_state: st.session_state.tidb_nodespec = list(tidb_spec_options.keys())[0] # Make sure current selection is valid if st.session_state.tidb_nodespec not in tidb_spec_options: st.session_state.tidb_nodespec = list(tidb_spec_options.keys())[0] tidb_nodespec = st.selectbox( "TiDB Node Spec", options=list(tidb_spec_options.keys()), format_func=lambda x: tidb_spec_options[x], index=list(tidb_spec_options.keys()).index( st.session_state.tidb_nodespec ), key="tidb_nodespec_selector", ) # Update session state st.session_state.tidb_nodespec = tidb_nodespec # Find vCPU count for selected TiDB spec for spec in tidb_specs: if spec["nodeSpecKey"] == tidb_nodespec: tidb_vcpu = spec["vCpu"] break else: st.warning(f"No TiDB node specs available for region {selected_region_id}") # Premium: RCU Configuration if plan_type == "Premium": st.subheader("Step 3: Capacity Configuration") col_rcu1, col_rcu2 = st.columns(2) with col_rcu1: premium_min_rcu = st.number_input( "Minimum Billing RCU", min_value=1000, value=5000, step=1000, help="Minimum Request Units billed per month" ) with col_rcu2: premium_max_rcu = st.number_input( "Maximum RCU", min_value=1000, value=20000, step=1000, help="Maximum Request Units allowed" ) # Step 4: Get available components based on TiDB selection (Dedicated only) if plan_type == "Dedicated": st.subheader("Step 4: Select Component Specifications") tikv_nodespec = None tiflash_nodespec = None pd_nodespec = None tiproxy_nodespec = None tikv_storage_type = None tiflash_storage_type = None constraints = {} tikv_specs = [] tiflash_specs = [] pd_specs = [] tiproxy_specs = [] if plan_type == "Dedicated" and selected_region_id and tidb_nodespec: all_nodespecs = cached_get_nodespecs(selected_region_id) available_comps = get_available_components(all_nodespecs, tidb_nodespec) tikv_specs = available_comps.get("tikv", []) tiflash_specs = available_comps.get("tiflash", []) pd_specs = available_comps.get("pd", []) tiproxy_specs = available_comps.get("tiproxy", []) constraints = available_comps.get("constraints", {}) # Show constraints info if there are special rules if tidb_vcpu in [2, 4]: constraint_msgs = [] if not constraints.get("tiflash_allowed"): constraint_msgs.append("TiFlash not supported") if not constraints.get("tiproxy_allowed"): constraint_msgs.append("TiProxy not supported") tidb_range = constraints.get("tidb_count_range", (1, None)) if tidb_range[1] is not None: constraint_msgs.append(f"TiDB nodes: {tidb_range[0]}-{tidb_range[1]}") tikv_range = constraints.get("tikv_count_range", (3, None)) if tikv_range[1] is not None: constraint_msgs.append(f"TiKV nodes: exactly {tikv_range[0]}") if constraint_msgs: st.info( f"**{tidb_vcpu}C Cluster Rules:** " + " | ".join(constraint_msgs) ) col_storage1, col_storage2 = st.columns(2) with col_storage1: st.markdown("### TiKV") if tikv_specs: tikv_spec_options = { s[ "nodeSpecKey" ]: f"{s['nodeSpecKey']} ({s['vCpu']}C{s['memorySizeGi']}G)" for s in tikv_specs } # Use session state for TiKV if "tikv_nodespec" not in st.session_state: st.session_state.tikv_nodespec = list(tikv_spec_options.keys())[0] if st.session_state.tikv_nodespec not in tikv_spec_options: st.session_state.tikv_nodespec = list(tikv_spec_options.keys())[0] tikv_nodespec = st.selectbox( "TiKV Node Spec", options=list(tikv_spec_options.keys()), format_func=lambda x: tikv_spec_options[x], index=list(tikv_spec_options.keys()).index( st.session_state.tikv_nodespec ), key="tikv_nodespec_selector", ) st.session_state.tikv_nodespec = tikv_nodespec # Get storage types for selected TiKV spec tikv_storage_types = [] for spec in tikv_specs: if spec["nodeSpecKey"] == tikv_nodespec: tikv_storage_types = spec.get("storageTypes", []) break if tikv_storage_types: # Set default storage type based on cloud provider default_tikv_storage = ( "Standard" if cloud_provider == "AWS" else "Basic" ) if default_tikv_storage not in tikv_storage_types: default_tikv_storage = tikv_storage_types[0] if "tikv_storage_type" not in st.session_state: st.session_state.tikv_storage_type = default_tikv_storage if st.session_state.tikv_storage_type not in tikv_storage_types: st.session_state.tikv_storage_type = tikv_storage_types[0] tikv_storage_type = st.selectbox( "TiKV Storage Type", options=tikv_storage_types, index=tikv_storage_types.index( st.session_state.tikv_storage_type ), key="tikv_storage_type_selector", ) st.session_state.tikv_storage_type = tikv_storage_type else: st.warning("No TiKV specs available") with col_storage2: st.markdown("### TiFlash") if constraints.get("tiflash_allowed", True) and tiflash_specs: tiflash_spec_options = { s[ "nodeSpecKey" ]: f"{s['nodeSpecKey']} ({s['vCpu']}C{s['memorySizeGi']}G)" for s in tiflash_specs } # Use session state for TiFlash if "tiflash_nodespec" not in st.session_state: st.session_state.tiflash_nodespec = list( tiflash_spec_options.keys() )[0] if st.session_state.tiflash_nodespec not in tiflash_spec_options: st.session_state.tiflash_nodespec = list( tiflash_spec_options.keys() )[0] tiflash_nodespec = st.selectbox( "TiFlash Node Spec", options=list(tiflash_spec_options.keys()), format_func=lambda x: tiflash_spec_options[x], index=list(tiflash_spec_options.keys()).index( st.session_state.tiflash_nodespec ), key="tiflash_nodespec_selector", ) st.session_state.tiflash_nodespec = tiflash_nodespec # Get storage types for selected TiFlash spec tiflash_storage_types = [] for spec in tiflash_specs: if spec["nodeSpecKey"] == tiflash_nodespec: tiflash_storage_types = spec.get("storageTypes", []) break if tiflash_storage_types: # TiFlash default is always Basic for all cloud providers default_tiflash_storage = "Basic" if default_tiflash_storage not in tiflash_storage_types: default_tiflash_storage = tiflash_storage_types[0] if "tiflash_storage_type" not in st.session_state: st.session_state.tiflash_storage_type = default_tiflash_storage if ( st.session_state.tiflash_storage_type not in tiflash_storage_types ): st.session_state.tiflash_storage_type = tiflash_storage_types[0] tiflash_storage_type = st.selectbox( "TiFlash Storage Type", options=tiflash_storage_types, index=tiflash_storage_types.index( st.session_state.tiflash_storage_type ), key="tiflash_storage_type_selector", ) st.session_state.tiflash_storage_type = tiflash_storage_type elif not constraints.get("tiflash_allowed", True): st.warning(f"TiFlash not supported for {tidb_vcpu}C clusters") else: st.info("No TiFlash specs available") # Workload Template st.subheader("Workload Template") # Template naming is not uniform: TPCC uses "-x-" while Sysbench/CH end # in "-x". Use the explicit plan registry instead of substring matching. available_templates = get_workload_templates_for_plan(plan_type) default_template = ( TPCC_1K_TEMPLATE if plan_type == "Dedicated" else TPCC_X_1K_TEMPLATE ) if "workload_template_choice" not in st.session_state: st.session_state.workload_template_choice = default_template else: st.session_state.workload_template_choice = ( canonicalize_workload_template_or_default( st.session_state.workload_template_choice, default=default_template ) if st.session_state.workload_template_choice in available_templates else default_template ) template_choice = st.selectbox( "Select Workload Template", available_templates, key="workload_template_choice", ) # Initialize workload config from the currently selected template on first # load if "workload_config_str" not in st.session_state: try: tpl = load_workload_template(template_choice) st.session_state.workload_config_str = json.dumps(tpl, indent=2) st.session_state.last_loaded_workload_template = template_choice except Exception as e: st.session_state.workload_config_str = '{"nodes": {}}' st.error(f"Failed to load template: {e}") else: # Automatically reload config when template selection changes last_tpl = st.session_state.get("last_loaded_workload_template") if last_tpl != template_choice: try: tpl = load_workload_template(template_choice) st.session_state.workload_config_str = json.dumps(tpl, indent=2) st.session_state.last_loaded_workload_template = template_choice except Exception as e: st.error(f"Failed to load template: {e}") # Default values for Premium (not used in Dedicated form but needed for json.loads) tidb_count = 1 tikv_count = 3 tikv_storage_size = 0 tiflash_count = 0 tiflash_storage_size = 0 tiproxy_count = 0 tidb_config = "{}" tikv_config = "{}" tiflash_config = "{}" tiproxy_config = "{}" with st.form("test_submission_form"): if plan_type == "Dedicated": st.subheader("Component Configuration") if plan_type == "Dedicated": # Display selected node specs as read-only info if tidb_nodespec: spec_info = f"**Selected:** TiDB: {tidb_nodespec}" if tikv_nodespec: spec_info += f" | TiKV: {tikv_nodespec}" if tiflash_nodespec: spec_info += f" | TiFlash: {tiflash_nodespec}" st.info(spec_info) # Get constraints for dynamic validation tidb_min = constraints.get("tidb_count_range", (1, None))[0] tidb_max = constraints.get("tidb_count_range", (1, None))[1] tikv_min = constraints.get("tikv_count_range", (3, None))[0] tikv_max = constraints.get("tikv_count_range", (3, None))[1] tikv_step = constraints.get("tikv_count_step", 3) col1, col2, col3, col4 = st.columns(4) with col1: st.markdown("### TiDB") st.text_input("TiDB SPU", value=tidb_nodespec or "", disabled=True) tidb_help = f"Min: {tidb_min}" if tidb_max: tidb_help += f", Max: {tidb_max}" tidb_count = st.number_input( "TiDB Count", min_value=tidb_min, max_value=tidb_max if tidb_max else 100, value=tidb_min, help=tidb_help, ) tidb_config = st.text_area("TiDB Config (JSON)", value="{}", height=100) with col2: st.markdown("### TiKV") st.text_input("TiKV SPU", value=tikv_nodespec or "", disabled=True) st.text_input( "Storage Type", value=tikv_storage_type or "", disabled=True, key="tikv_storage_type_display", ) tikv_storage_size = st.number_input( "Storage Size (Gi)", min_value=0, value=1000, help="Storage size in GiB", key="tikv_storage_size", ) tikv_help = f"Min: {tikv_min}, Must be multiple of {tikv_step}" if tikv_max: tikv_help = f"Must be exactly {tikv_min}" tikv_count = st.number_input( "TiKV Count", min_value=tikv_min, max_value=tikv_max if tikv_max else 100, value=tikv_min, step=tikv_step, help=tikv_help, ) tikv_config = st.text_area("TiKV Config (JSON)", value="{}", height=100) with col3: st.markdown("### TiFlash") tiflash_disabled = not constraints.get("tiflash_allowed", True) st.text_input("TiFlash SPU", value=tiflash_nodespec or "", disabled=True) st.text_input( "Storage Type", value=tiflash_storage_type or "", disabled=True, key="tiflash_storage_type_display", ) tiflash_storage_size = st.number_input( "Storage Size (Gi)", min_value=0, value=1000, disabled=tiflash_disabled, help="Storage size in GiB", key="tiflash_storage_size", ) tiflash_count = st.number_input( "TiFlash Count", min_value=0, value=0 if tiflash_disabled else 1, disabled=tiflash_disabled, help=("Not supported for 2C/4C clusters" if tiflash_disabled else None), ) tiflash_config = st.text_area( "TiFlash Config (JSON)", value="{}", height=100, disabled=tiflash_disabled, ) with col4: st.markdown("### TiProxy") tiproxy_disabled = not constraints.get("tiproxy_allowed", True) if not tiproxy_disabled and tiproxy_specs and len(tiproxy_specs) > 0: tiproxy_spec_options = { s[ "nodeSpecKey" ]: f"{s['nodeSpecKey']} ({s['vCpu']}C{s['memorySizeGi']}G)" for s in tiproxy_specs } # Use session state for TiProxy if "tiproxy_nodespec" not in st.session_state: st.session_state.tiproxy_nodespec = list( tiproxy_spec_options.keys() )[0] if st.session_state.tiproxy_nodespec not in tiproxy_spec_options: st.session_state.tiproxy_nodespec = list( tiproxy_spec_options.keys() )[0] tiproxy_nodespec = st.selectbox( "TiProxy Node Spec", options=list(tiproxy_spec_options.keys()), format_func=lambda x: tiproxy_spec_options[x], index=list(tiproxy_spec_options.keys()).index( st.session_state.tiproxy_nodespec ), key="tiproxy_nodespec_selector", ) st.session_state.tiproxy_nodespec = tiproxy_nodespec else: if tiproxy_disabled: st.warning("TiProxy not supported for 2C clusters") else: st.info("No TiProxy specs available") tiproxy_nodespec = None tiproxy_count = st.number_input( "TiProxy Count", min_value=0, value=0, disabled=tiproxy_disabled, help=( "Not supported for 2C clusters" if tiproxy_disabled else "Valid: 0, 2, 3, 6, 9, 12, 15, 18, 21, 24" ), ) tiproxy_config = st.text_area( "TiProxy Config (JSON)", value="{}", height=100, disabled=tiproxy_disabled, ) # Use default values when PD configuration is disabled pd_nodespec = None pd_count = 3 if plan_type == "Dedicated" else 1 # Premium: 1 PD pd_config = "{}" # Default empty config # Workload configuration section below component configuration st.subheader("Advanced Options") # Different UI for bench-ch template if template_choice in ("bench-ch", "bench-ch-x"): is_premium_ch = template_choice == "bench-ch-x" with st.expander("CH Workload Parameters"): st.caption("Configure CH benchmark parameters") # Parse current config try: current_config = json.loads(st.session_state.workload_config_str) params = current_config.get("params", {}) except: params = {} # CH-specific input fields ch_threads = st.text_input( "Threads", value=str(params.get("threads", "100")), key="ch_threads", ) ch_duration = st.text_input( "Duration", value=params.get("duration", "2h"), help="e.g., 20m, 4h", key="ch_duration", ) ch_warehouses = st.text_input( "Warehouses", value=str(params.get("warehouses", "1000")), key="ch_warehouses", ) ch_workload_memory = st.text_input( "Workload Node Memory", value=params.get("workload_node_memory", "16Gi"), help="e.g., 16Gi, 32Gi", key="ch_workload_memory", ) ch_workload_cpu = st.text_input( "Workload Node CPU(k8s core)", value=params.get("workload_node_cpu", "8000m"), help="e.g., 8000m, 16000m", key="ch_workload_cpu", ) # Update the workload_config_str with the CH params ch_config = { "template_type": "bench-ch", "params": { "threads": ch_threads, "duration": ch_duration, "warehouses": ch_warehouses, "workload_node_memory": ch_workload_memory, "workload_node_cpu": ch_workload_cpu, }, } st.session_state.workload_config_str = json.dumps(ch_config, indent=2) else: # Original workload config for other templates with st.expander("Workload Config (JSON)"): st.caption("DON'T MODIFY IT IF UNCERTAIN!") st.text_area( "Workload Config (JSON)", key="workload_config_str", height=300, ) submitted = st.form_submit_button("Submit Test", type="primary") if submitted: # Validation errors = [] if not user_email or not selected_region_id: errors.append( "Missing required fields: Email, Cloud Provider, Region" ) if plan_type == "Dedicated" and not tidb_nodespec: errors.append("Missing required field: TiDB Node Spec") # Validate org_id/project_id inputs (optional, but must be provided together) if bool(org_id_input) ^ bool(project_id_input): errors.append( "Please set both org_id and project_id, or leave both blank." ) # Validate email format - must contain @pingcap. if user_email and "@pingcap." not in user_email: errors.append( "Email must be a PingCAP email address (must contain '@pingcap.')" ) # Validate TiKV count (Dedicated only) if plan_type == "Dedicated": if tikv_count % tikv_step != 0: errors.append(f"TiKV count must be multiple of {tikv_step}") # Validate TiProxy count (only specific values allowed) valid_tiproxy_counts = [0, 2, 3, 6, 9, 12, 15, 18, 21, 24] if tiproxy_count not in valid_tiproxy_counts: errors.append(f"TiProxy count must be one of: {valid_tiproxy_counts}") # Premium: validate RCU if plan_type == "Premium": if premium_min_rcu > premium_max_rcu: errors.append("Minimum RCU must be <= Maximum RCU") # Validate TiFlash for 2C/4C clusters (Dedicated only) if plan_type == "Dedicated" and not constraints.get("tiflash_allowed", True) and tiflash_count > 0: errors.append(f"TiFlash not supported for {tidb_vcpu}C clusters") # Validate TiProxy for 2C clusters (Dedicated only) if plan_type == "Dedicated" and not constraints.get("tiproxy_allowed", True) and tiproxy_count > 0: errors.append(f"TiProxy not supported for {tidb_vcpu}C clusters") if errors: for error in errors: st.error(error) else: try: # Parse JSON configs tidb_cfg = json.loads(tidb_config) tikv_cfg = json.loads(tikv_config) tiflash_cfg = json.loads(tiflash_config) pd_cfg = json.loads(pd_config) tiproxy_cfg = json.loads(tiproxy_config) workload_cfg = json.loads(st.session_state.workload_config_str) task_data = { "user_email": user_email, "cloud_provider": cloud_provider, "region": selected_region_id, "cluster_version": cluster_version, "org_id": org_id_input or None, "project_id": project_id_input or None, "tidb_spu": tidb_nodespec if plan_type == "Dedicated" else "", "tidb_count": int(tidb_count) if plan_type == "Dedicated" else 1, "tidb_config": tidb_cfg if plan_type == "Dedicated" else {}, "tikv_spu": (tikv_nodespec or "") if plan_type == "Dedicated" else "", "tikv_count": int(tikv_count) if plan_type == "Dedicated" else 3, "tikv_config": tikv_cfg if plan_type == "Dedicated" else {}, "tikv_storage_type": ( tikv_storage_type.lower() if tikv_storage_type else None ) if plan_type == "Dedicated" else None, "tikv_storage_size_gi": int(tikv_storage_size) if plan_type == "Dedicated" else 0, "tiflash_spu": (tiflash_nodespec or "") if plan_type == "Dedicated" else "", "tiflash_count": int(tiflash_count) if plan_type == "Dedicated" else 0, "tiflash_config": tiflash_cfg if plan_type == "Dedicated" else {}, "tiflash_storage_type": ( tiflash_storage_type.lower() if tiflash_storage_type else None ) if plan_type == "Dedicated" else None, "tiflash_storage_size_gi": ( int(tiflash_storage_size) if tiflash_count > 0 else 0 ) if plan_type == "Dedicated" else 0, "pd_spu": pd_nodespec or "", "pd_count": int(pd_count), "pd_config": pd_cfg, "tiproxy_spu": tiproxy_nodespec or "", "tiproxy_count": int(tiproxy_count) if plan_type == "Dedicated" else 0, "tiproxy_config": tiproxy_cfg if plan_type == "Dedicated" else {}, "workload_template": template_choice, "workload_config": workload_cfg, } # Premium: add RCU to workload_config if plan_type == "Premium": task_data["workload_config"] = build_premium_workload_config( workload_cfg, min_rcu=premium_min_rcu, max_rcu=premium_max_rcu, ) task_id = db.create_task(task_data) if task_id: st.success(f"Task submitted successfully! Task ID: {task_id}") st.balloons() else: st.error( "Failed to submit task. Please check database connection." ) except json.JSONDecodeError as e: st.error(f"Invalid JSON in configuration: {e}")