ci4sManagement-cloud/react-ui/src/pages/Experiment/index.jsx

import CommonTableCell from '@/components/CommonTableCell';
import KFIcon from '@/components/KFIcon';
import { ExperimentStatus, TensorBoardStatus } from '@/enums';
import {
  deleteExperimentById,
  getExperiment,
  getExperimentById,
  getQueryByExperimentId,
  getTensorBoardStatusReq,
  postExperiment,
  putExperiment,
  runExperiments,
  runTensorBoardReq,
} from '@/services/experiment/index.js';
import { getWorkflow } from '@/services/pipeline/index.js';
import themes from '@/styles/theme.less';
import { to } from '@/utils/promise';
import { modalConfirm } from '@/utils/ui';
import { App, Button, ConfigProvider, Dropdown, Space, Table } from 'antd';
import classNames from 'classnames';
import { useEffect, useRef, useState } from 'react';
import { useNavigate } from 'react-router-dom';
import { ComparisonType } from './Comparison/config';
import AddExperimentModal from './components/AddExperimentModal';
import ExperimentInstance from './components/ExperimentInstance';
import Styles from './index.less';
import { experimentStatusInfo } from './status';

// 定时器
const timerIds = new Map();

function Experiment() {
  const navgite = useNavigate();
  const [experimentList, setExperimentList] = useState([]);
  const [workflowList, setWorkflowList] = useState([]);
  const [queryFlow, setQueryFlow] = useState({
    offset: 1,
    page: 0,
    size: 10000,
    name: null,
  });
  const [experimentId, setExperimentId] = useState(null);
  const [experimentInList, setExperimentInList] = useState([]);
  const [expandedRowKeys, setExpandedRowKeys] = useState(null);
  const [total, setTotal] = useState(0);
  const [isAdd, setIsAdd] = useState(true);
  const [isModalOpen, setIsModalOpen] = useState(false);
  const [addFormData, setAddFormData] = useState({});
  const [experimentInsTotal, setExperimentInsTotal] = useState(0);
  const { message } = App.useApp();
  const pageOption = useRef({ page: 1, size: 10 });
  const paginationProps = {
    showSizeChanger: true,
    showQuickJumper: true,
    showTotal: () => `共${total}条`,
    total: total,
    page: pageOption.current.page,
    size: pageOption.current.size,
    onChange: (current, size) => paginationChange(current, size),
  };

  useEffect(() => {
    getList();
    getWorkflowList();
    return () => {
      clearExperimentInTimers();
    };
  }, []);

  // 获取实验列表
  const getList = async () => {
    const params = {
      offset: 0,
      page: pageOption.current.page - 1,
      size: pageOption.current.size,
    };
    const [res, _] = await to(getExperiment(params));
    if (res && res.data && Array.isArray(res.data.content)) {
      setExperimentList(
        res.data.content.map((item) => {
          return { ...item, key: item.id };
        }),
      );

      setTotal(res.data.totalElements);
    }
  };

  // 获取流水线列表
  const getWorkflowList = async () => {
    const [res, _] = await to(getWorkflow(queryFlow));
    if (res && res.data && res.data.content) {
      setWorkflowList(res.data.content);
    }
  };

  // 获取实验实例列表
  const getQueryByExperiment = async (experimentId, page) => {
    const params = {
      experimentId: experimentId,
      page: page,
      size: 5,
    };
    const [res, error] = await to(getQueryByExperimentId(params));
    if (res && res.data) {
      const { content = [], totalElements = 0 } = res.data;
      setExpandedRowKeys(experimentId);
      try {
        const list = content.map((v) => {
          const nodes_result = v.nodes_result ? JSON.parse(v.nodes_result) : {};
          return {
            ...v,
            nodes_result,
          };
        });
        if (page === 0) {
          setExperimentInList(list);
          clearExperimentInTimers();
        } else {
          setExperimentInList((prev) => [...prev, ...list]);
        }
        setExperimentInsTotal(totalElements);
        // 获取 TensorBoard 状态
        list.forEach((item) => {
          if (item.nodes_result?.tensorboard_log) {
            getTensorBoardStatus(item);
          }
        });
      } catch (error) {
        console.error('JSON parse error: ', error);
      }
    }
  };

  // 运行 TensorBoard
  const runTensorBoard = async (experimentIn) => {
    const params = {
      namespace: experimentIn.nodes_result.tensorboard_log.namespace,
      path: experimentIn.nodes_result.tensorboard_log.path,
      pvc_name: experimentIn.nodes_result.tensorboard_log.pvc_name,
    };
    const [res] = await to(runTensorBoardReq(params));
    if (res) {
      experimentIn.tensorboardUrl = res.data;
      const timerId = timerIds.get(experimentIn.id);
      if (timerId) {
        clearTimeout(timerId);
        timerIds.delete(experimentIn.id);
        getTensorBoardStatus(experimentIn);
      }
    }
  };

  // 获取 TensorBoard 状态
  const getTensorBoardStatus = async (experimentIn) => {
    const params = {
      namespace: experimentIn.nodes_result.tensorboard_log.namespace,
      path: experimentIn.nodes_result.tensorboard_log.path,
      pvc_name: experimentIn.nodes_result.tensorboard_log.pvc_name,
    };
    const [res] = await to(getTensorBoardStatusReq(params));
    if (res && res.data) {
      setExperimentInList((prevList) => {
        return prevList.map((item) => {
          if (item.id === experimentIn.id) {
            return {
              ...item,
              tensorBoardStatus: res.data.status,
              tensorboardUrl: res.data.url,
            };
          }
          return item;
        });
      });

      let timerId = timerIds.get(experimentIn.id);
      if (timerId) {
        clearTimeout(timerId);
        timerIds.delete(experimentIn.id);
      }
      timerId = setTimeout(() => {
        getTensorBoardStatus(experimentIn);
      }, 10 * 1000);
      timerIds.set(experimentIn.id, timerId);
    }
  };

  // 展开实例
  const expandChange = (e, record) => {
    clearExperimentInTimers();
    setExperimentInList([]);
    if (record.id === expandedRowKeys) {
      setExpandedRowKeys(null);
    } else {
      getQueryByExperiment(record.id, 0);
    }
  };

  // 终止实验实例获取 TensorBoard 状态的定时器
  const clearExperimentInTimers = () => {
    timerIds.values().forEach((timerId) => {
      clearTimeout(timerId);
    });
    timerIds.clear();
  };

  // 创建实验
  const createExperiment = () => {
    setIsAdd(true);
    setAddFormData({});
    setExperimentId(null);
    setIsModalOpen(true);
  };

  // 编辑实验
  const editExperiment = (id) => {
    getExperimentById(id).then((res) => {
      setAddFormData({
        ...res.data,
      });
      setExperimentId(res.data.id);
      setIsAdd(false);
      setIsModalOpen(true);
    });
  };
  // 创建或编辑实验取消
  const handleCancel = () => {
    setIsModalOpen(false);
  };

  // 创建或者编辑实验接口请求
  const handleAddExperiment = async (values) => {
    const global_param = JSON.stringify(values.global_param);
    if (!experimentId) {
      const params = {
        ...values,
        global_param,
      };
      const [res, _] = await to(postExperiment(params));
      if (res) {
        message.success('新建实验成功');
        setIsModalOpen(false);
        getList();
      }
    } else {
      const params = { ...values, global_param, id: experimentId };
      const [res, _] = await to(putExperiment(params));
      if (res) {
        message.success('编辑实验成功');
        setIsModalOpen(false);
        getList();
      }
    }
  };

  // 当前页面切换
  const paginationChange = async (current, size) => {
    pageOption.current = {
      page: current,
      size: size,
    };
    getList();
  };
  // 运行实验
  const runExperiment = async (id) => {
    const [res] = await to(runExperiments(id));
    if (res) {
      message.success('运行成功');
      refreshExperimentIns(id);
    } else {
      message.error('运行失败');
    }
  };

  // 跳转到流水线
  const gotoPipeline = (e, record) => {
    e.stopPropagation();
    navgite({ pathname: `/pipeline/template/${record.workflow_id}` });
  };

  // 跳转到实验实例详情
  const gotoInstanceInfo = (item, record) => {
    navgite({ pathname: `/pipeline/experiment/${record.workflow_id}/${item.id}` });
  };

  // 处理 TensorBoard 操作
  const handleTensorboard = async (experimentIn) => {
    if (
      experimentIn.tensorBoardStatus === TensorBoardStatus.Terminated ||
      experimentIn.tensorBoardStatus === TensorBoardStatus.Failed
    ) {
      await runTensorBoard(experimentIn);
    } else if (
      experimentIn.tensorBoardStatus === TensorBoardStatus.Running &&
      experimentIn.tensorboardUrl
    ) {
      window.open(experimentIn.tensorboardUrl, '_blank');
    }
  };

  // 实验实例终止
  const handleInstanceTerminate = async (experimentIn) => {
    setExperimentInList((prevList) => {
      return prevList.map((item) => {
        if (item.id === experimentIn.id) {
          return {
            ...item,
            status: ExperimentStatus.Terminated,
          };
        }
        return item;
      });
    });
  };

  // 实验对比菜单
  const getComparisonMenu = (experimentId) => {
    return {
      items: [
        {
          label: <span>训练对比</span>,
          key: ComparisonType.Train,
        },
        {
          label: <span>评估对比</span>,
          key: ComparisonType.Evaluate,
        },
      ],
      onClick: ({ key }) => {
        navgite(`/pipeline/experiment/compare?type=${key}&id=${experimentId}`);
      },
    };
  };

  // 刷新实验实例列表
  const refreshExperimentIns = (experimentId) => {
    getQueryByExperiment(experimentId, 0);
  };

  // 加载更多实验实例
  const loadMoreExperimentIns = () => {
    const page = Math.round(experimentInList.length / 5);
    getQueryByExperiment(expandedRowKeys, page);
  };

  const columns = [
    {
      title: '实验名称',
      dataIndex: 'name',
      key: 'name',
      render: (text) => <div>{text}</div>,
      width: '16%',
    },
    {
      title: '关联流水线名称',
      dataIndex: 'workflow_name',
      key: 'workflow_name',
      render: (text, record) => <a onClick={(e) => gotoPipeline(e, record)}>{text}</a>,
      width: '16%',
    },
    {
      title: '实验描述',
      dataIndex: 'description',
      key: 'description',
      render: CommonTableCell(true),
      ellipsis: { showTitle: false },
    },
    {
      title: '最近五次运行状态',
      dataIndex: 'status_list',
      key: 'status_list',
      width: 200,
      render: (text) => {
        let newText = text && text.replace(/\s+/g, '').split(',');
        return (
          <>
            {newText && newText.length > 0
              ? newText.map((item, index) => {
                  return (
                    <img
                      style={{ width: '17px', marginRight: '6px' }}
                      key={index}
                      src={experimentStatusInfo[item].icon}
                    />
                  );
                })
              : null}
          </>
        );
      },
    },

    {
      title: '操作',
      key: 'action',
      width: 350,
      render: (_, record) => (
        <Space size="small">
          <Button
            type="link"
            size="small"
            key="run"
            icon={<KFIcon type="icon-yunhang" />}
            onClick={() => {
              runExperiment(record.id);
            }}
          >
            运行
          </Button>
          <Button
            type="link"
            size="small"
            key="edit"
            icon={<KFIcon type="icon-bianji" />}
            onClick={() => {
              editExperiment(record.id);
            }}
          >
            编辑
          </Button>
          <Dropdown key="comparison" menu={getComparisonMenu(record.id)}>
            <a onClick={(e) => e.preventDefault()}>
              <Space style={{ padding: '0 7px' }}>
                <KFIcon type="icon-shiyanduibi" />
                实验对比
              </Space>
            </a>
          </Dropdown>
          <ConfigProvider
            theme={{
              token: {
                colorLink: themes['warningColor'],
              },
            }}
          >
            <Button
              type="link"
              size="small"
              key="batchRemove"
              icon={<KFIcon type="icon-shanchu" />}
              onClick={() => {
                modalConfirm({
                  title: '删除后，该实验将不可恢复',
                  content: '是否确认删除?',
                  onOk: () => {
                    deleteExperimentById(record.id).then((ret) => {
                      if (ret.code === 200) {
                        message.success('删除成功');
                        getList();
                      } else {
                        message.error(ret.msg);
                      }
                    });
                  },
                });
              }}
            >
              删除
            </Button>
          </ConfigProvider>
        </Space>
      ),
    },
  ];
  return (
    <div className={Styles.experimentBox}>
      <div className={Styles.experimentTopBox}>
        <Button type="default" onClick={createExperiment} icon={<KFIcon type="icon-xinjian2" />}>
          新建实验
        </Button>
      </div>
      <div className={classNames('vertical-scroll-table', Styles.experimentTable)}>
        <Table
          columns={columns}
          dataSource={experimentList}
          pagination={paginationProps}
          rowKey="id"
          scroll={{ y: 'calc(100% - 55px)' }}
          expandable={{
            expandedRowRender: (record) => (
              <ExperimentInstance
                experimentInList={experimentInList}
                experimentInsTotal={experimentInsTotal}
                onClickInstance={(item) => gotoInstanceInfo(item, record)}
                onClickTensorBoard={handleTensorboard}
                onRemove={() => refreshExperimentIns(record.id)}
                onTerminate={handleInstanceTerminate}
                onLoadMore={() => loadMoreExperimentIns()}
              ></ExperimentInstance>
            ),
            onExpand: (e, a) => {
              expandChange(e, a);
            },
            expandedRowKeys: [expandedRowKeys],
            rowExpandable: (record) => true,
          }}
        />
      </div>

      {isModalOpen && (
        <AddExperimentModal
          isAdd={isAdd}
          open={isModalOpen}
          initialValues={addFormData}
          onCancel={handleCancel}
          onFinish={handleAddExperiment}
          workflowList={workflowList}
        />
      )}
    </div>
  );
}
export default Experiment;