ygo 1.0.2__py3-none-any.whl → 1.2.12__py3-none-any.whl

ygo/__init__.py

@@ -1,10 +1,31 @@
-# -*- coding: utf-8 -*-
-"""
----------------------------------------------
-Created on 2025/4/28 15:25
-@author: ZhangYundi
-@email: yundi.xxii@outlook.com
----------------------------------------------
-"""
+from .exceptions import FailTaskError, WarnException
+from .pool import pool
+from .delay import delay
+from .utils import (
+    fn_params,
+    fn_signature_params,
+    fn_path,
+    fn_code,
+    fn_info,
+    module_from_str,
+    fn_from_str,
+    locate,
+)
+from .lazy import lazy_import
 
-from .ygo import *
+__version__ = "1.2.12"
+
+__all__ = [
+    "FailTaskError",
+    "delay",
+    "WarnException",
+    "fn_params",
+    "fn_signature_params",
+    "fn_path",
+    "fn_code",
+    "fn_info",
+    "fn_from_str",
+    "module_from_str",
+    "pool",
+    "lazy_import"
+]

ygo/delay.py

@@ -0,0 +1,89 @@
+# -*- coding: utf-8 -*-
+"""
+---------------------------------------------
+Created on 2025/5/26 20:12
+@author: ZhangYundi
+@email: yundi.xxii@outlook.com
+---------------------------------------------
+"""
+
+import functools
+import inspect
+
+
+class DelayedFunction:
+
+    def __init__(self, func):
+        self.func = func
+        self._fn_params_k = inspect.signature(self.func).parameters.keys()
+        self.stored_kwargs = self._get_default_args(func)
+        if hasattr(func, 'stored_kwargs'):
+            self.stored_kwargs.update(func.stored_kwargs)
+
+    def _get_default_args(self, func):
+        signature = inspect.signature(func)
+        return {
+            k: v.default
+            for k, v in signature.parameters.items()
+            if v.default is not inspect.Parameter.empty
+        }
+
+    def __call__(self, *args, **kwargs):
+        def delayed(*args, **_kwargs):
+            new_kwargs = {k: v for k, v in self.stored_kwargs.items()}
+            for k, v in _kwargs.items():
+                if k not in self._fn_params_k:
+                    continue
+                new_kwargs[k] = v
+            return self.func(*args, **new_kwargs)
+
+        self._stored_kwargs(**kwargs)
+        new_fn = functools.wraps(self.func)(delayed)
+        new_fn.stored_kwargs = self.stored_kwargs
+        return new_fn
+
+    def _stored_kwargs(self, **kwargs):
+        for k, v in kwargs.items():
+            if k not in self._fn_params_k:
+                continue
+            self.stored_kwargs[k] = v
+
+
+def delay(func):
+    """
+    延迟执行
+    Parameters
+    ----------
+    func: Callable
+        需要延迟执行的对象, 必须是一个Callable对象
+
+    Returns
+    -------
+    DelayedFunction
+        将预先设置好的参数包装进原始的Callable对象中
+
+    Examples
+    --------
+
+    场景1：基本使用
+
+    >>> fn = delay(lambda a, b: a+b)(a=1, b=2)
+    >>> fn()
+    3
+
+    场景2: 逐步传递参数
+
+    >>> fn1 = delay(lambda a, b, c: a+b+c)(a=1)
+    >>> fn2 = delay(fn1)(b=2)
+    >>> fn2(c=3)
+    6
+
+    场景3: 参数更改
+
+    >>> fn1 = delay(lambda a, b, c: a+b+c)(a=1, b=2)
+    >>> fn2 = delay(fn1)(c=3, b=5)
+    >>> fn2()
+    9
+    """
+    return DelayedFunction(func)
+

ygo/exceptions.py

@@ -7,7 +7,23 @@ Created on 2024/12/18 下午7:01
 ---------------------------------------------
 """
 
+from dataclasses import dataclass
+
 class WarnException(Exception):
     """自定义异常类，仅用于警告"""
     def __init__(self, message):
-        super().__init__(message)  # 调用父类的构造函数
+        super().__init__(message)  # 调用父类的构造函数
+
+@dataclass
+class FailTaskError:
+    task_name: str
+    error: Exception
+
+    def __str__(self):
+        return f"""
+[失败任务]: {self.task_name}
+[错误信息]: \n{self.error}
+"""
+
+    def __repr__(self):
+        return self.__str__()

ygo/lazy.py

@@ -0,0 +1,50 @@
+# -*- coding: utf-8 -*-
+"""
+---------------------------------------------
+Copyright (c) 2025 ZhangYundi
+Licensed under the MIT License. 
+Created on 2025/6/29 10:36
+Email: yundi.xxii@outlook.com
+Description:
+---------------------------------------------
+"""
+
+from typing import Any
+from .utils import locate
+
+class LazyImport:
+    def __init__(self, full_name: str):
+        self._full_name = full_name
+        self._obj = None
+
+    def _load(self) -> Any:
+        if self._obj is None:
+            self._obj = locate(self._full_name)
+        return self._obj
+
+    def __getattr__(self, attr: str) -> Any:
+        obj = self._load()
+        return getattr(obj, attr)
+
+    def __dir__(self) -> list[str]:
+        obj = self._load()
+        return dir(obj)
+
+    def __call__(self, *args, **kwargs) -> Any:
+        obj = self._load()
+        if isinstance(obj, type):
+            return obj(*args, **kwargs)
+        elif callable(obj):
+            return obj(*args, **kwargs)
+        else:
+            raise TypeError(f"The target `{self._full_name}` is not callable or instantiable.")
+
+    def __str__(self):
+        return self._full_name
+
+    def __repr__(self):
+        return self._full_name
+
+def lazy_import(full_name: str):
+    """实现模块和方法的懒加载"""
+    return LazyImport(full_name)

ygo/pool.py

@@ -0,0 +1,286 @@
+# -*- coding: utf-8 -*-
+"""
+---------------------------------------------
+Created on 2024/11/4 下午2:10
+@author: ZhangYundi
+@email: yundi.xxii@outlook.com
+---------------------------------------------
+"""
+import functools
+import multiprocessing
+import os
+import sys
+import threading
+import warnings
+from typing import Literal
+
+import logair
+from joblib import Parallel, delayed
+
+from .delay import delay
+from .exceptions import WarnException, FailTaskError
+
+with warnings.catch_warnings():
+    warnings.simplefilter("ignore")
+    from tqdm.auto import tqdm
+
+logger = logair.get_logger("ygo")
+
+
+def run_job(job, task_id, queue):
+    """执行任务并更新队列"""
+    try:
+        result = job()
+    except WarnException as e:
+        logger.warning(FailTaskError(task_name=job.task_name, error=e))
+        result = None
+    except Exception as e:
+        logger.error(FailTaskError(task_name=job.task_name, error=e), exc_info=e)
+        result = None
+    queue.put((task_id, 1))
+    return result
+
+
+def update_progress_bars(tqdm_objects: list[tqdm],
+                         task_ids,
+                         queue: multiprocessing.Queue,
+                         num_tasks: int,
+                         task_counts: dict):
+    """根据队列中的消息更新 tqdm 进度条"""
+    completed_tasks = 0
+    completed_task_jobs = {id_: 0 for id_ in task_ids}
+    while completed_tasks < num_tasks:
+        task_id, progress_value = queue.get()  # 从队列获取进度更新
+        completed_task_jobs[task_id] += 1
+        if completed_task_jobs[task_id] >= task_counts[task_id]:
+            completed_tasks += 1
+        tqdm_objects[task_id].update(progress_value)
+    [tqdm_object.close() for tqdm_object in tqdm_objects]
+
+
+class pool:
+    """
+    每个fn运行一次算一个job，每个job需要指定job_name, 如果没有job_name, 则默认分配 TaskDefault
+    相同 job_name 的fn归到同一个task, 同时该task命名为job_name
+    即一个task中包含了多个需要运行的 job fn
+    task1 <job_fn1, job_fn2, ...>
+    task2 <job_fn3, job_fn4, ...>
+    所有的job_fn都会通过joblib并行运行
+    """
+
+    def __init__(self,
+                 n_jobs: int = -1,
+                 show_progress: bool = True,
+                 backend: Literal['loky', 'threading', 'multiprocessing'] = 'loky',
+                 ):
+        """backend: loky/threading/multiprocessing"""
+        self.show_progress = show_progress
+        if n_jobs < 0:
+            n_jobs = multiprocessing.cpu_count() - 1  # 给系统留一个cpu核
+        self._n_jobs = n_jobs
+
+        default_kwargs = {
+            # 'verbose': 10,
+            'batch_size': 'auto',
+            'pre_dispatch': f"2*{n_jobs}",
+            'max_nbytes': '1M',  # 减少进程间通信
+            'timeout': None,
+            # 'prefer': 'processes',
+        }
+        self._parallel = Parallel(n_jobs=self._n_jobs, verbose=0, backend=backend,
+                                  **default_kwargs) if self._n_jobs > 1 else None
+        self._default_task_name = "GO-JOB"
+        self._all_jobs = list()  # list[job]
+        self._all_tasks = list()  # list[task_name]
+        self._task_ids = dict()  # {task_name1: 0, task_name2: 1, ...}
+        self._task_counts = dict()
+        self._leave_mapping = dict()
+        self._id_counts = dict()
+        self._manager = None
+
+        # 1. 配置环境
+        self._configure_environment()
+
+        # 2. 抢占CPU核心
+        if self._n_jobs > 1:
+            self._steal_cpu_cores()
+
+        # 3. 预分配内存（如果启用）
+        # if self.memory_aware:
+        #     self._preallocate_memory_pool()
+
+    def _configure_environment(self):
+        """配置环境以最大化资源使用"""
+        # 禁用所有竞争对手的并行
+        env_vars = {
+            'OMP_NUM_THREADS': '1',
+            'MKL_NUM_THREADS': '1',
+            'OPENBLAS_NUM_THREADS': '1',
+            'VECLIB_MAXIMUM_THREADS': '1',
+            'NUMEXPR_NUM_THREADS': '1',
+            'NUMBA_NUM_THREADS': '1',
+            'PYTHON_GIL': '0',  # 尝试禁用GIL（如果支持）
+        }
+
+        for key, value in env_vars.items():
+            os.environ[key] = value
+
+        # 设置Python优化
+        sys.setrecursionlimit(1000000)
+
+        # 设置joblib参数
+        os.environ['LOKY_MAX_CPU_COUNT'] = str(self._n_jobs)
+        os.environ['JOBLIB_START_METHOD'] = 'forkserver'  # 更快的进程启动
+
+    def _steal_cpu_cores(self):
+        """为worker进程抢占CPU核心"""
+        try:
+            import psutil
+
+            # 分析CPU使用情况
+            cpu_percent = psutil.cpu_percent(interval=0.2, percpu=True)
+
+            # 分配核心给worker
+            worker_cores = []
+            for i in range(self._n_jobs):
+                # 寻找负载最低的核心
+                available_cores = [j for j in range(len(cpu_percent))
+                                   if j not in worker_cores]
+
+                if available_cores:
+                    # 选择负载最低的可用核心
+                    available_loads = [(j, cpu_percent[j]) for j in available_cores]
+                    available_loads.sort(key=lambda x: x[1])
+                    chosen_core = available_loads[0][0]
+                    worker_cores.append(chosen_core)
+
+            self.worker_affinity = {
+                i: [core] for i, core in enumerate(worker_cores)
+            }
+
+            # print(f"🎯 为 {self.n_jobs} 个worker分配核心: {worker_cores}")
+
+        except Exception as e:
+            print(f"⚠️ CPU核心分配失败: {e}")
+
+    def submit(self, fn, job_name: str = "", postfix: str = "", leave: bool = True):
+        """
+        提交并行任务
+        Parameters
+        ----------
+        fn: callable
+            需要并行的callable对象
+        job_name: str
+            提交到的任务名, 不同的任务对应不同的进度条
+        postfix: str
+            后缀
+        leave: bool
+            进度条完成后是否保留在屏幕上，默认 True
+        Returns
+        -------
+
+        Examples
+        --------
+        import time
+        import ygo
+        >>> def task_fn1(a, b):
+                time.sleep(3)
+                return a+b
+        >>> def task_fn2():
+                time.sleep(5)
+                return 0
+        >>> with ygo.pool() as go:
+                go.submit(task_fn1, job_name="task1")(a=1, b=2)
+                go.submit(task_fn2, job_name="task2")()
+
+                go.do()
+        """
+
+        # 提交任务，对任务进行分类，提交到对应的task id中，并且封装新的功能：使其在运行完毕后将任务进度更新放入队列
+        @functools.wraps(fn)
+        def collect(**kwargs):
+            """归集所有的job到对应的task"""
+            with warnings.catch_warnings():
+                warnings.simplefilter('ignore')
+                job = delay(fn)(**kwargs)
+                task_name = self._default_task_name if not job_name else job_name
+                task_name = f"{task_name}::{postfix}"
+                if task_name not in self._task_ids:
+                    self._task_ids[task_name] = len(self._all_tasks)
+                    self._task_counts[task_name] = 0
+                    self._all_tasks.append(task_name)
+                    self._id_counts[self._task_ids[task_name]] = 0
+                    self._leave_mapping[task_name] = leave
+                self._task_counts[task_name] += 1
+                self._id_counts[self._task_ids[task_name]] += 1
+                job.task_id = self._task_ids[task_name]
+                job.job_id = self._task_counts[task_name]
+                job.task_name = task_name
+                self._all_jobs.append(job)
+                return job
+
+        return collect
+
+    def do(self):
+        if self.show_progress:
+            # 消息队列进行通信
+            self._manager = multiprocessing.Manager()
+            queue = self._manager.Queue()
+            tqdm_bars = [tqdm(total=self._task_counts[task_name],
+                              desc=f"{str(task_name.split('::')[0])}",
+                              postfix=task_name.split("::")[1],
+                              leave=self._leave_mapping.get(task_name, True)) for task_name in
+                         self._all_tasks]
+            # 初始化多个任务的进度条，每个任务一个
+            task_ids = [task_id for task_id in range(len(self._all_tasks))]
+            # 创建并启动用于更新进度条的线程
+            progress_thread = threading.Thread(target=update_progress_bars, args=(
+                tqdm_bars, task_ids, queue, len(self._all_tasks), self._id_counts))
+            progress_thread.start()
+            if self._parallel is not None:
+                with warnings.catch_warnings():
+                    warnings.filterwarnings("ignore", category=UserWarning)
+                    # 执行并行任务
+                    result = self._parallel(delayed(run_job)(job=job,
+                                                             task_id=job.task_id,
+                                                             queue=queue) for job in self._all_jobs)
+            else:
+                result = [run_job(job=job, task_id=job.task_id, queue=queue) for job in self._all_jobs]
+            # 等待进度更新线程执行完毕
+            progress_thread.join()
+        else:
+            if self._parallel is not None:
+                with warnings.catch_warnings():
+                    warnings.filterwarnings("ignore", category=UserWarning)
+                    result = self._parallel(delayed(job)() for job in self._all_jobs)
+            else:
+                result = [job() for job in self._all_jobs]
+        self._all_jobs = list()  # list[job]
+        self._all_tasks = list()  # list[task_name]
+        self._task_ids = dict()  # {task_name1: 0, task_name2: 1, ...}
+        self._task_counts = dict()
+        self._id_counts = dict()
+        self._leave_mapping = dict()
+        self.close()
+        return result
+
+    def close(self):
+        """释放所有资源"""
+        if hasattr(self, '_parallel') and self._parallel is not None:
+            try:
+                self._parallel.__exit__(None, None, None)
+            except Exception as e:
+                logger.warning(f"Failed to close Parallel: {e}")
+
+        if hasattr(self, '_manager') and self._manager is not None:
+            try:
+                self._manager.shutdown()
+            except Exception as e:
+                logger.warning(f"Failed to shutdown Manager: {e}")
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        # 释放进程
+        self.close()