smartlibs 0.1.9__py3-none-any.whl

auto_tasks/aaas/__utils.py

@@ -0,0 +1,6 @@
+import logging
+
+from smart.auto import AutoLoad
+
+logger = logging.getLogger('auto_tasks')
+auto_load = AutoLoad()

auto_tasks/aaas/aaas_task.py

@@ -0,0 +1,62 @@
+from smart.auto.tree import TreeMultiTask
+from smart.aaas.client import AaasClient
+from smart.utils.yaml import yaml_dumps
+
+from .__utils import auto_load, logger
+
+
+@auto_load.task('aaas__client')
+class AaasTask(TreeMultiTask):
+    CTX_TASK_LIST_NAME = 'aaas__client:tasks'
+
+    def conn(self, entrypoint:str=None, namespace=None, module=None, enable_https=False):
+        client = AaasClient(
+            entrypoint = entrypoint,
+            namespace = namespace,
+            enable_https = enable_https
+        )
+
+        if module:
+            client.set_module(module)
+
+        return {
+            'client': client
+        }
+    
+    def run(self, client:AaasClient, task_name, task_module=None, task_id=None, \
+            task_configs=None, bind_arg=None, run_opts=None, state_hook=None):
+
+        create_rst = client.create_task(
+            task_name=task_name,
+            task_id=task_id,
+            module=task_module,
+            configs=task_configs,
+            bind_arg=bind_arg,
+            run_opts=run_opts,
+            state_hook=state_hook
+        )
+
+        task_id = create_rst.get('task_id')
+
+        if task_id:
+            tasks = self.context.list(self.CTX_TASK_LIST_NAME)
+            tasks.append({
+                'client': client.init_args(),
+                'task_id': task_id,
+            })
+
+        logger.info('aaas__client create_task: %s', create_rst)
+    
+    def asdl(self, client:AaasClient, task_module=None, task_configs=None, bind_arg=None, run_opts=None):
+        asdl_rst = client.asdl(
+            module=task_module,
+            configs=task_configs,
+            bind_arg=bind_arg,
+            run_opts=run_opts,
+        )
+
+        logger.info('aaas__client asdl_rst:\n%s', yaml_dumps(asdl_rst.get("result") or asdl_rst))
+
+        return {
+            "asdl": asdl_rst
+        }

auto_tasks/gpu/GpuTool.py

@@ -0,0 +1,82 @@
+from .__utils import logger
+from collections import namedtuple
+import random
+
+
+MemInfo = namedtuple('MemInfo', ['total', 'free', 'used'])
+
+
+class GpuInfoGetter:
+    def get_device_count(self) -> int:
+        """获取显卡数量
+
+        Returns:
+            int: 显卡数量
+        """
+        pass
+
+    def get_memory_info(self, index:int) -> MemInfo:
+        """获取指定显卡的显存信息
+
+        Args:
+            index (int): 显卡序号, 0表示第一块显卡
+
+        Returns:
+            MemInfo: 显存信息
+        """
+        pass
+
+    def find_mem_free_device(self, gpu_num:int=1, shuffle:bool=False,
+                free_memory:int=None, free_memory_ratio:float=None, filter_fn:callable=None)->list:
+        """查找显存足够的显卡列表
+        free_memory和free_memory_ratio都为None时, 所有显卡都可返回。
+        shuffle=True可随机获取可用显卡；shuffle=False则按顺序检查可用显卡。
+        查找可用显卡以执行代码的时刻的显存来判断，与实际占用显卡一般有一段间隔。
+        在多个任务并发时，可能会出现多个任务同时在抢占同一张显卡。shuffle=True能减少抢占情况。
+
+        Args:
+            gpu_num (int, optional): 需要查找的设备数量. Defaults to 1.
+            shuffle (bool, optional): 是否随机打乱GPU获取顺序. Defaults to False.
+            free_memory (int, optional): 过滤可用显存小于free_memory的显卡. Defaults to None.
+            free_memory_ratio (float, optional): 过滤可用显存/显存小于free_memory_ratio的显卡. Defaults to None.
+            filter_fn (callable, optional): 过滤可用显卡的函数, 例如lambda idx, memInfo:True将过滤所有显卡. Defaults to None.
+
+        Returns:
+            list: 数据结构为[tuple(显卡序号:int, 显存信息:MemInfo)]
+        """
+        gpu_count = self.get_device_count()
+        gpu_index_iter = range(gpu_count)
+
+        if shuffle:
+            gpu_index_iter = list(gpu_index_iter)
+            random.shuffle(gpu_index_iter)
+
+        choosed_gpu = []
+
+        for gpu_index in gpu_index_iter:
+            meminfo:MemInfo = self.get_memory_info(gpu_index)
+
+            _choose = True
+            try:
+                if free_memory_ratio is not None:
+                    _ratio = float(meminfo.free) / meminfo.total
+                    if _ratio < free_memory_ratio:
+                        _choose = False
+                
+                if free_memory is not None:
+                    if meminfo.free < free_memory:
+                        _choose = False
+                
+                if filter_fn and filter_fn(gpu_index, meminfo):
+                    _choose = False
+            except Exception as err:
+                logger.warning("find_mem_free_device err: %s", err)
+                _choose = False
+            
+            if _choose:
+                choosed_gpu.append((gpu_index, meminfo))
+            
+            if len(choosed_gpu) >= gpu_num:
+                break
+        
+        return choosed_gpu

auto_tasks/gpu/__utils.py

@@ -0,0 +1,8 @@
+import logging
+
+from smart.auto import AutoLoad
+
+logger = logging.getLogger('auto_tasks')
+auto_load = AutoLoad()
+
+task_hook = auto_load.hook

auto_tasks/gpu/cuda_tool.py

@@ -0,0 +1,204 @@
+import os, time
+
+from smart.auto import TreeMultiTask
+from smart.utils.env import auto_set_env_by_prefix, AppEnv
+from smart.utils.cast import cast_bool
+from .__utils import auto_load, logger, task_hook
+from .GpuTool import GpuInfoGetter, MemInfo
+
+
+_options = {
+    "nvml_is_disable": False
+}
+
+try:
+    import pynvml
+except:
+    _options['nvml_is_disable'] = True
+    _options['err_msg'] = "miss pynvml"
+
+
+class CudaInfoGetter(GpuInfoGetter):
+    def __init__(self) -> None:
+        self.__inited = False
+    
+    def __nvmlInit(self):
+        if not self.__inited:
+            pynvml.nvmlInit()
+            self.__inited = True
+
+    def get_device_count(self):
+        self.__nvmlInit()
+        return pynvml.nvmlDeviceGetCount()
+        
+    def get_memory_info(self, index:int):
+        self.__nvmlInit()
+        handle = pynvml.nvmlDeviceGetHandleByIndex(index)
+        meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
+        return MemInfo(
+            meminfo.total, meminfo.free, meminfo.used
+        )
+
+
+@auto_load.task('gpu_tools.cuda_tool')
+class CudaToolTask(TreeMultiTask):
+    def find_available_gpu(self, gpu_num:int=1, shuffle:bool=False,
+            free_memory:int=None, free_memory_ratio:float=None, 
+            device_env_key:str='CUDA_VISIBLE_DEVICES', ctx_state_name:str='cuda_tool'):
+        """查找可用的GPU
+        当free_memory和free_memory_ratio都为None时，所有显卡都为可用，返回的available_gpu数组长度=min(gpu_num, 机器实际显卡数)。
+        shuffle=True可随机获取可用显卡；shuffle=False则按顺序检查可用显卡。
+        查找可用显卡以执行代码的时刻的显存来判断，与实际占用显卡一般有一段间隔。
+        在多个任务并发时，可能会出现多个任务同时在抢占同一张显卡。shuffle=True能减少抢占情况。
+
+        Args:
+            gpu_num (int, optional): 需要的GPU数量. Defaults to 1.
+            shuffle (bool, optional): 是否随机打乱GPU顺序. Defaults to False.
+            free_memory (int, optional): 显卡的可用显存小于free_memory为不可用. Defaults to None.
+            free_memory_ratio (float, optional): 显卡的可用显存/显存小于free_memory_ratio为不可用. Defaults to None.
+            device_env_key (str, optional): 将可用的设备序号设置到环境变量中, 空值表示不设置环境变量. Defaults to 'CUDA_VISIBLE_DEVICES'.
+            ctx_state_name (str, optional): 将返回的available_gpu列表保存到context中. Defaults to 'cuda_tool'.
+
+        Returns:
+            dict: {"available_gpu":[(显卡序号:int, 显存信息:MemInfo)]}
+        """
+        if not cast_bool(AppEnv.get("USE_GPU", True)):
+            return
+
+        if _options['nvml_is_disable']:
+            logger.error("nvml is not support. %s", _options.get("err_msg", ""))
+            return
+
+        cuda = CudaInfoGetter()
+
+        with self.context.store.lock((ctx_state_name, "find_available_gpu")):
+            used_gpu = self.context.list((ctx_state_name, "used_gpu"))
+            used_gpu_idx = [idx for idx, _ in used_gpu]
+
+            choosed_gpu = cuda.find_mem_free_device(
+                gpu_num=gpu_num, 
+                shuffle=shuffle,
+                free_memory=free_memory, 
+                free_memory_ratio=free_memory_ratio,
+                filter_fn=lambda idx, _:(idx in used_gpu_idx)
+            )
+
+            if choosed_gpu:
+                self.context.list((ctx_state_name, "used_gpu")).extend(choosed_gpu)
+        
+        if device_env_key:
+            device_env_val = ','.join([
+                str(val[0]) for val in choosed_gpu
+            ])
+            auto_set_env_by_prefix(device_env_key, device_env_val)
+            logger.info("cuda_tool.find_available_gpu set_env %s: %s", device_env_key, device_env_val)
+
+        return {
+            "available_gpu": choosed_gpu
+        }
+    
+    @task_hook.before_task()
+    def hook_available_gpu(self, gpu_num:int=1, shuffle:bool=False,
+            free_memory:int=None, free_memory_ratio:float=None, 
+            device_env_key:str=None, ctx_state_name:str='cuda_tool'):
+        """查找可用的GPU的勾子函数(在其他任务启动前先执行)
+        当free_memory和free_memory_ratio都为None时，所有显卡都为可用，返回的available_gpu数组长度=min(gpu_num, 机器实际显卡数)。
+        shuffle=True可随机获取可用显卡；shuffle=False则按顺序检查可用显卡。
+        查找可用显卡以执行代码的时刻的显存来判断，与实际占用显卡一般有一段间隔。
+        在多个任务并发时，可能会出现多个任务同时在抢占同一张显卡。shuffle=True能减少抢占情况。
+
+        Args:
+            gpu_num (int, optional): 需要的GPU数量. Defaults to 1.
+            shuffle (bool, optional): 是否随机打乱GPU顺序. Defaults to False.
+            free_memory (int, optional): 显卡的可用显存小于free_memory为不可用. Defaults to None.
+            free_memory_ratio (float, optional): 显卡的可用显存/显存小于free_memory_ratio为不可用. Defaults to None.
+            device_env_key (str, optional): 将可用的设备序号设置到环境变量中, 空值表示不设置环境变量. Defaults to None.
+            ctx_state_name (str, optional): 将返回的available_gpu列表保存到context中. Defaults to 'cuda_tool'.
+
+        Returns:
+            dict: {"available_gpu":[(显卡序号:int, 显存信息:MemInfo)]}
+        """
+        if not cast_bool(AppEnv.get("USE_GPU", True)):
+            return
+        if _options['nvml_is_disable']:
+            logger.error("nvml is not support. %s", _options.get("err_msg", ""))
+            return
+
+        cuda = CudaInfoGetter()
+
+        choosed_gpu = cuda.find_mem_free_device(
+            gpu_num=gpu_num, 
+            shuffle=shuffle,
+            free_memory=free_memory, 
+            free_memory_ratio=free_memory_ratio
+        )
+
+        if device_env_key:
+            device_env_val = ','.join([
+                str(val[0]) for val in choosed_gpu
+            ])
+            auto_set_env_by_prefix(device_env_key, device_env_val)
+            logger.info("cuda_tool.hook_available_gpu set_env %s: %s", device_env_key, device_env_val)
+
+        if ctx_state_name:
+            self.context.state(ctx_state_name).update({
+                "available_gpu_num": len(choosed_gpu),
+                "available_gpu_list": choosed_gpu
+            })
+
+        return {
+            "available_gpu": choosed_gpu
+        }
+
+    def pop_gpu_from_ctx(self, gpu_num:int=1, min_gpu_num:int=0, 
+                device_env_key:str='CUDA_VISIBLE_DEVICES', ctx_state_name:str='cuda_tool'):
+        """从context中获取可用GPU, 并将GPU序号保存到环境变量
+        本方法与hook_available_gpu搭配使用. 
+        当可用的GPU数量少于min_gpu_num时, 则不从context的available_gpu_list中pop数据, 同时CUDA_VISIBLE_DEVICES环境变量设置为-1. 
+
+        Args:
+            gpu_num (int, optional): 需要的gpu数量. Defaults to 1.
+            min_gpu_num (int, optional): 最小需要的gpu数量. Defaults to 0.
+            device_env_key (str, optional): 保存GPU序号的环境变量, 本参数一般不修改. Defaults to 'CUDA_VISIBLE_DEVICES'.
+            ctx_state_name (str, optional): 保存可用GPU列表的context名称, 本参数一般不修改. Defaults to 'cuda_tool'.
+
+        Returns:
+            dict: {available_gpu:[(显卡序号:int, 显存信息:MemInfo)]}
+        """
+        if not cast_bool(AppEnv.get("USE_GPU", True)):
+            return
+            
+        if gpu_num < min_gpu_num:
+            gpu_num = min_gpu_num
+
+        ctx_state = self.context.state(ctx_state_name)
+        available_gpu_num = ctx_state.wait("available_gpu_num")
+        choosed_gpu = []
+
+        if available_gpu_num >= min_gpu_num:
+            with self.context.store.lock((ctx_state_name, "pop_gpu_from_ctx")):
+                gpu_list = ctx_state.get("available_gpu_list")
+                logger.debug("CudaToolTask.pop_gpu_from_ctx current gpu_list=%s", gpu_list)
+                # time.sleep(1)
+
+                try:
+                    if gpu_list and len(gpu_list) >= min_gpu_num:
+                        for i in range(gpu_num):
+                            gpu_idx_info_tuple = gpu_list.pop(0)
+                            choosed_gpu.append(gpu_idx_info_tuple)
+                except IndexError:
+                    logger.info("CudaToolTask.pop_gpu_from_ctx no enough gpu")
+                
+                ctx_state.set("available_gpu_list", gpu_list)
+        
+        if device_env_key:
+            os.environ[device_env_key] = ", ".join(
+                str(i) for i, _ in choosed_gpu
+            ) if len(choosed_gpu) else "-1"
+            logger.info("CudaToolTask.pop_gpu_from_ctx set %s=%s", device_env_key, os.environ[device_env_key])
+        
+        return {
+            "available_gpu": choosed_gpu
+        }
+    
+

auto_tasks/jsonl/__utils.py

@@ -0,0 +1,6 @@
+import logging
+
+from smart.auto import AutoLoad
+
+logger = logging.getLogger('auto_tasks')
+auto_load = AutoLoad()

auto_tasks/jsonl/jsonl_file.py

@@ -0,0 +1,132 @@
+import os, json, logging
+
+from smart.auto import TreeMultiTask, AutoLoad
+from smart.utils import list_safe_iter, path_join
+
+from .__utils import auto_load, logger
+
+
+@auto_load.task('jsonl__file')
+class JsonlFileTask(TreeMultiTask):
+    def __resolve_file_name_keys(self, file_name_keys, file_name_idx_or_key=None):
+        if isinstance(file_name_keys, str):
+            file_name_key_list = [k.strip() for k in file_name_keys.split(',')]
+        else:
+            file_name_key_list = file_name_keys
+            
+        if isinstance(file_name_idx_or_key, int):
+            yield file_name_key_list[file_name_idx_or_key]
+        elif file_name_idx_or_key:
+            yield file_name_idx_or_key
+        else:
+            yield from file_name_key_list
+
+    def pattern_read(self, file_name_keys, file_name_pattern:dict='{}', dir_path=None, file_open_opts=None, root_dir=None, group_key='_group', file_path=None):
+        """读取多份jsonl文件
+
+        Arguments:
+            file_name_keys {list} -- 文件名的键列表
+
+        Keyword Arguments:
+            file_name_pattern {str} -- 文件名模版, 使用'{}'占位file_name_key (default: {'{}'})
+            dir_path {str} -- 文件目录路径 (default: {None})
+            file_open_opts {dict} -- 打开文件选项 (default: {None})
+            root_dir {str} -- 根路径 (default: {None})
+            group_key {str} -- item的分组键, 值为file_name_key (default: {'_group'})
+            file_path {str} -- 弃用, 请使用dir_path代替 (default: {None})
+
+        Returns:
+            dict -- {item_iter_fn}
+        """
+        _file_name_keys = file_name_keys
+        dir_path = dir_path or file_path
+        file_open_opts = {'mode': 'r', 'encoding': 'utf8', **(file_open_opts or {})}
+
+        def item_iter_fn(file_name_idx_or_key = None, file_name_keys=None):
+            file_name_key_list = self.__resolve_file_name_keys(file_name_keys or _file_name_keys, file_name_idx_or_key)
+            
+            for file_name_key in file_name_key_list:
+                num_items = 0
+                file_name = file_name_pattern.format(file_name_key)
+                file = path_join(root_dir, dir_path, file_name)
+
+                if not os.path.exists(file):
+                    logger.warning('JsonlFileTask.pattern_read: no found file %s', file)
+                    continue
+
+                with open(file, **file_open_opts) as f:
+                    for line in f:
+                        if not line:
+                            continue
+                        item = json.loads(line)
+                        if isinstance(item, dict) and group_key:
+                            item[group_key] = file_name_key
+                        yield item
+                        num_items += 1
+                logger.debug('JsonlFileTask.pattern_read %s items from %s', num_items, file)
+                
+        return {
+            'item_iter_fn': item_iter_fn
+        }
+    
+    def read(self, file_name, dir_path=None, file_open_opts:dict=None, root_dir=None, file_path=None):
+        """读取jsonl文件
+
+        Arguments:
+            file_name {str} -- 文件名
+
+        Keyword Arguments:
+            dir_path {str} -- 文件目录路径 (default: {None})
+            file_open_opts {dict} -- 打开文件选项 (default: {None})
+            root_dir {str} -- 根路径 (default: {None})
+            file_path {str} -- 弃用, 请使用dir_path代替 (default: {None})
+
+        Returns:
+            dict -- {item_iter_fn}
+        """
+        dir_path = dir_path or file_path
+        file = path_join(root_dir, dir_path, file_name)
+        file_open_opts = {'mode': 'r', 'encoding': 'utf8', **(file_open_opts or {})}
+
+        def item_iter_fn():
+            with open(file, **file_open_opts) as f:
+                for line in f:
+                    item = json.loads(line)
+                    yield item
+        
+        return {
+            'item_iter_fn': item_iter_fn
+        }
+    
+    def write(self, file_name, dir_path=None, file_open_opts:dict=None, root_dir=None, item_iter=None, item_iter_fn=None, recv_args={}, file_path=None):
+        """写jsonl文件
+
+        Arguments:
+            file_name {str} -- 文件名
+
+        Keyword Arguments:
+            dir_path {str} -- 文件目录路径 (default: {None})
+            file_open_opts {dict} -- 打开文件选项 (default: {None})
+            root_dir {str} -- 根路径 (default: {None})
+            item_iter {generator} -- item生成器 (default: {None})
+            item_iter_fn {callable} -- item生成器构造函数; item_iter非空时, 本参数无效 (default: {None})
+            recv_args {dict} -- 接收数据函数的参数选项; item_iter非空时, 本参数无效 (default: {{}})
+            file_path {str} -- 弃用, 请使用dir_path代替 (default: {None})
+        """
+        assert file_name
+        dir_path = dir_path or file_path
+        file = path_join(root_dir, dir_path, file_name, auto_mkdir=True)
+        file_open_opts = {'mode': 'w', 'encoding': 'utf8', **(file_open_opts or {})}
+        logger.info('jsonl__file.write %s', file_name)
+
+        item_iter = item_iter or (item_iter_fn or self.recv_data)(**recv_args)
+
+        count = 0
+        with open(file, **file_open_opts) as f:
+            for item in item_iter:
+                json.dump(item, f, ensure_ascii=False)
+                f.write('\n')
+                count += 1
+        
+        logger.debug('jsonl__file.write %s items to %s', count, file_name)
+

auto_tasks/lib/__utils.py

@@ -0,0 +1,6 @@
+import logging
+
+from smart.auto import AutoLoad
+
+logger = logging.getLogger('auto_tasks')
+auto_load = AutoLoad()

auto_tasks/lib/auto.py

@@ -0,0 +1,26 @@
+from smart.auto.tree import TreeMultiTask
+from smart.auto.Runner import AutoRunner
+from smart.auto.ctx.runner_context import WithAutoRunner
+
+from .__utils import auto_load, logger
+
+
+@auto_load.task('lib.auto')
+class LibAuto(TreeMultiTask, WithAutoRunner):
+    def run_tree(self, tree_name):
+        auto_runner:AutoRunner = getattr(self, 'auto_runner', None)
+        logger.debug('LibAuto.run_tree: %s', tree_name)
+
+        if auto_runner:
+            auto_runner.start(tree_name)
+        else:
+            logger.error('!!! lib.auto.run_tree fail because auto_runner is empty, maybe you should use lib.auto as task_cls')
+
+    def run_task(self, task_exp):
+        auto_runner:AutoRunner = getattr(self, 'auto_runner', None)
+        logger.debug('LibAuto.run_task: %s', task_exp)
+
+        if auto_runner:
+            auto_runner.start(task_exp, default_ns='task')
+        else:
+            logger.error('!!! lib.auto.run_task fail because auto_runner is empty, maybe you should use lib.auto as task_cls')