Before, we only decide whether to collect data by current step,
it will not work well in dataset sink mode, so we check to see
if it's a dataset sink mode, and decide whether to collect data.