from __future__ import annotations

from typing import Any, Optional

class DatasetFactory:
    def __init__(self) -> None: ...
    def create_dataset(self, datafeed_class: str = ...): ...

class DatasetBase:
    proto_desc: Any = ...
    dataset: Any = ...
    thread_num: int = ...
    filelist: Any = ...
    use_ps_gpu: bool = ...
    psgpu: Any = ...
    def __init__(self) -> None: ...
    def set_pipe_command(self, pipe_command: Any) -> None: ...
    def set_so_parser_name(self, so_parser_name: Any) -> None: ...
    def set_rank_offset(self, rank_offset: Any) -> None: ...
    fea_eval: Any = ...
    def set_fea_eval(self, record_candidate_size: Any, fea_eval: bool = ...) -> None: ...
    def slots_shuffle(self, slots: Any) -> None: ...
    def set_batch_size(self, batch_size: Any) -> None: ...
    def set_pv_batch_size(self, pv_batch_size: Any) -> None: ...
    def set_thread(self, thread_num: Any) -> None: ...
    def set_filelist(self, filelist: Any) -> None: ...
    def set_input_type(self, input_type: Any) -> None: ...
    def set_use_var(self, var_list: Any) -> None: ...
    def set_hdfs_config(self, fs_name: Any, fs_ugi: Any) -> None: ...
    def set_download_cmd(self, download_cmd: Any) -> None: ...
    def desc(self): ...

class InMemoryDataset(DatasetBase):
    fleet_send_batch_size: Any = ...
    is_user_set_queue_num: bool = ...
    queue_num: Any = ...
    parse_ins_id: bool = ...
    parse_content: bool = ...
    parse_logkey: bool = ...
    merge_by_sid: bool = ...
    enable_pv_merge: bool = ...
    merge_by_lineid: bool = ...
    fleet_send_sleep_seconds: Any = ...
    trainer_num: int = ...
    def __init__(self) -> None: ...
    dataset: Any = ...
    def set_feed_type(self, data_feed_type: Any) -> None: ...
    def set_queue_num(self, queue_num: Any) -> None: ...
    def set_parse_ins_id(self, parse_ins_id: Any) -> None: ...
    def set_parse_content(self, parse_content: Any) -> None: ...
    def set_parse_logkey(self, parse_logkey: Any) -> None: ...
    def set_merge_by_sid(self, merge_by_sid: Any) -> None: ...
    def set_enable_pv_merge(self, enable_pv_merge: Any) -> None: ...
    def preprocess_instance(self) -> None: ...
    def set_current_phase(self, current_phase: Any) -> None: ...
    def postprocess_instance(self) -> None: ...
    def set_fleet_send_batch_size(self, fleet_send_batch_size: int = ...) -> None: ...
    def set_fleet_send_sleep_seconds(self, fleet_send_sleep_seconds: int = ...) -> None: ...
    def set_merge_by_lineid(self, merge_size: int = ...) -> None: ...
    gen_uni_feasigns: Any = ...
    local_shard_num: Any = ...
    def set_generate_unique_feasigns(self, generate_uni_feasigns: Any, shard_num: Any) -> None: ...
    def generate_local_tables_unlock(
        self, table_id: Any, fea_dim: Any, read_thread_num: Any, consume_thread_num: Any, shard_num: Any
    ) -> None: ...
    def set_date(self, date: Any) -> None: ...
    def load_into_memory(self, is_shuffle: bool = ...) -> None: ...
    def preload_into_memory(self, thread_num: Optional[Any] = ...) -> None: ...
    def wait_preload_done(self) -> None: ...
    def local_shuffle(self) -> None: ...
    def global_shuffle(self, fleet: Optional[Any] = ..., thread_num: int = ...) -> None: ...
    def release_memory(self) -> None: ...
    def get_pv_data_size(self): ...
    def get_memory_data_size(self, fleet: Optional[Any] = ...): ...
    def get_shuffle_data_size(self, fleet: Optional[Any] = ...): ...

class QueueDataset(DatasetBase):
    def __init__(self) -> None: ...
    def local_shuffle(self) -> None: ...
    def global_shuffle(self, fleet: Optional[Any] = ...) -> None: ...

class FileInstantDataset(DatasetBase):
    def __init__(self) -> None: ...
    def local_shuffle(self) -> None: ...
    def global_shuffle(self, fleet: Optional[Any] = ...) -> None: ...

class BoxPSDataset(InMemoryDataset):
    boxps: Any = ...
    def __init__(self) -> None: ...
    def set_date(self, date: Any) -> None: ...
    def begin_pass(self) -> None: ...
    def end_pass(self, need_save_delta: Any) -> None: ...
    def wait_preload_done(self) -> None: ...
    def load_into_memory(self) -> None: ...
    def preload_into_memory(self) -> None: ...
    def slots_shuffle(self, slots: Any) -> None: ...
