# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved.

# Copyright 2025 The TransferQueue Team

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

import os

from .client import TransferQueueClient

from .dataloader import StreamingDataLoader, StreamingDataset

from .interface import (

async_kv_batch_get,

async_kv_batch_get_by_meta,

async_kv_batch_put,

async_kv_clear,

async_kv_list,

async_kv_put,

close,

get_client,

get_metrics_endpoint,

init,

kv_batch_get,

kv_batch_get_by_meta,

kv_batch_put,

kv_clear,

kv_list,

kv_put,

)

from .metadata import BatchMeta, KVBatchMeta

from .sampler import BaseSampler

from .sampler.grpo_group_n_sampler import GRPOGroupNSampler

from .sampler.rank_aware_sampler import RankAwareSampler

from .sampler.seqlen_balanced_sampler import SeqlenBalancedSampler

from .sampler.sequential_sampler import SequentialSampler

__all__ = (

[

# High-Level KV Interface

"init",

"close",

"get_metrics_endpoint",

"kv_put",

"kv_batch_put",

"kv_batch_get",

"kv_batch_get_by_meta",

"kv_list",

"kv_clear",

"async_kv_put",

"async_kv_batch_put",

"async_kv_batch_get",

"async_kv_batch_get_by_meta",

"async_kv_list",

"async_kv_clear",

"KVBatchMeta",

]

+ [

# High-Level StreamingDataLoader Interface

"StreamingDataset",

"StreamingDataLoader",

]

+ [

# Low-Level Native Interface

"get_client",

"BatchMeta",

"TransferQueueClient",

]

+ [

# Sampler

"BaseSampler",

"GRPOGroupNSampler",

"SequentialSampler",

"RankAwareSampler",

"SeqlenBalancedSampler",

]

)

version_folder = os.path.dirname(os.path.join(os.path.abspath(__file__)))

with open(os.path.join(version_folder, "version/version")) as f:

__version__ = f.read().strip()