ci4sManagement-cloud/react-ui/src/pages/Experiment/Comparison/index.tsx

/*
 * @Author: 赵伟
 * @Date: 2024-10-10 09:55:12
 * @Description: 实验对比
 */

import TableColTitle from '@/components/TableColTitle';
import {
  getExpEvaluateInfosReq,
  getExpMetricsReq,
  getExpTrainInfosReq,
} from '@/services/experiment';
import { tableSorter } from '@/utils';
import { to } from '@/utils/promise';
import SessionStorage from '@/utils/sessionStorage';
import tableCellRender, { TableCellValueType } from '@/utils/table';
import { useNavigate, useSearchParams } from '@umijs/max';
import { App, Button, Table, TablePaginationConfig, TableProps } from 'antd';
import classNames from 'classnames';
import { useEffect, useMemo, useState } from 'react';
import ExperimentStatusCell from '../components/ExperimentStatusCell';
import { ComparisonType, comparisonConfig } from './config';
import styles from './index.less';

type TableData = {
  experiment_ins_id: number;
  run_id: string;
  dataset: string[];
  start_time: string;
  status: string;
  metrics_names?: string[];
  metrics?: Record<string, number>;
  params_names?: string[];
  params?: Record<string, number>;
};

function ExperimentComparison() {
  const [searchParams] = useSearchParams();
  const comparisonType = searchParams.get('type') as ComparisonType;
  const experimentId = searchParams.get('id');
  const [tableData, setTableData] = useState<TableData[]>([]);
  const [selectedRowKeys, setSelectedRowKeys] = useState<React.Key[]>([]);
  const [total, setTotal] = useState(0);
  const [pagination, setPagination] = useState<TablePaginationConfig>({
    current: 1,
    pageSize: 10,
  });

  const { message } = App.useApp();
  const navigate = useNavigate();
  const config = comparisonConfig[comparisonType];

  useEffect(() => {
    // 获取对比数据列表
    const getComparisonData = async () => {
      const request =
        comparisonType === ComparisonType.Train ? getExpTrainInfosReq : getExpEvaluateInfosReq;
      const params = {
        page: pagination.current! - 1,
        size: pagination.pageSize,
      };
      const [res] = await to(request(experimentId, params));
      if (res && res.data) {
        const { content = [], totalElements = 0 } = res.data;
        setTableData(content);
        setTotal(totalElements);
      }
    };

    getComparisonData();
  }, [experimentId, pagination, comparisonType]);

  // 获取对比 url
  const getExpMetrics = async () => {
    const [res] = await to(getExpMetricsReq(selectedRowKeys));
    if (res && res.data) {
      const url = res.data;
      // window.open(url, '_blank');
      SessionStorage.setItem(SessionStorage.aimUrlKey, url);
      navigate('compare-visual');
    }
  };

  // 对比按钮 click
  const handleComparisonClick = () => {
    if (selectedRowKeys.length < 2) {
      message.error('请至少选择两项进行对比');
      return;
    }
    getExpMetrics();
  };

  // 选择行
  const rowSelection: TableProps<TableData>['rowSelection'] = {
    type: 'checkbox',
    columnWidth: 48,
    fixed: 'left',
    selectedRowKeys,
    onChange: (selectedRowKeys: React.Key[]) => {
      setSelectedRowKeys(selectedRowKeys);
    },
  };

  // 分页切换
  const handleTableChange: TableProps<TableData>['onChange'] = (
    pagination,
    _filters,
    _sorter,
    { action },
  ) => {
    if (action === 'paginate') {
      setPagination(pagination);
    }
  };

  const columns: TableProps<TableData>['columns'] = useMemo(() => {
    const first: TableData | undefined = tableData.find(
      (item) => item.metrics_names && item.metrics_names.length > 0,
    );
    const metricsNames = first?.metrics_names ?? [];
    const paramsNames = first?.params_names ?? [];
    return [
      {
        title: '基本信息',
        align: 'center',
        children: [
          {
            title: '实例 ID',
            dataIndex: 'experiment_ins_id',
            key: 'experiment_ins_id',
            width: 100,
            fixed: 'left',
            align: 'center',
            render: tableCellRender(),
          },
          {
            title: '运行时间',
            dataIndex: 'start_time',
            key: 'start_time',
            width: 180,
            fixed: 'left',
            align: 'center',
            render: tableCellRender(false, TableCellValueType.Date),
          },
          {
            title: '运行状态',
            dataIndex: 'status',
            key: 'status',
            width: 100,
            fixed: 'left',
            // align: 'center',
            render: ExperimentStatusCell,
          },
          {
            title: `${config.title}数据集`,
            dataIndex: 'dataset',
            key: 'dataset',
            width: 180,
            fixed: 'left',
            align: 'center',
            render: tableCellRender(true, TableCellValueType.Array),
          },
        ],
      },
      {
        title: `${config.title}参数`,
        align: 'center',
        children: paramsNames.map((name) => ({
          title: <TableColTitle title={name} />,
          dataIndex: ['params', name],
          key: name,
          width: 150,
          align: 'center',
          render: tableCellRender(true),
          sorter: (a, b) => tableSorter(a.params?.[name], b.params?.[name]),
          showSorterTooltip: false,
        })),
      },
      {
        title: `${config.title}指标`,
        align: 'center',
        children: metricsNames.map((name) => ({
          title: <TableColTitle title={name} />,
          dataIndex: ['metrics', name],
          key: name,
          width: 150,
          align: 'center',
          render: tableCellRender(true),
          sorter: (a, b) => tableSorter(a.metrics?.[name], b.metrics?.[name]),
          showSorterTooltip: false,
        })),
      },
    ];
  }, [tableData, config]);

  return (
    <div className={styles['experiment-comparison']}>
      <div className={styles['experiment-comparison__header']}>
        <Button type="default" onClick={handleComparisonClick}>
          可视化对比
        </Button>
      </div>
      <div className={classNames('vertical-scroll-table', styles['experiment-comparison__table'])}>
        <Table
          dataSource={tableData}
          columns={columns}
          rowSelection={rowSelection}
          scroll={{ y: 'calc(100% - 110px)', x: '100%' }}
          pagination={{
            ...pagination,
            total: total,
            showSizeChanger: true,
            showQuickJumper: true,
          }}
          onChange={handleTableChange}
          bordered={true}
          rowKey={(record) => record.run_id || record.experiment_ins_id}
        />
      </div>
    </div>
  );
}

export default ExperimentComparison;