speedy-utils 1.1.5__py3-none-any.whl → 1.1.7__py3-none-any.whl

speedy_utils/common/notebook_utils.py

@@ -8,12 +8,12 @@ from IPython.display import HTML, display
 from tabulate import tabulate
 
 
-def change_dir(target_directory: str = 'POLY') -> None:
+def change_dir(target_directory: str = "POLY") -> None:
     """Change directory to the first occurrence of x in the current path."""
-    cur_dir = pathlib.Path('./')
+    cur_dir = pathlib.Path("./")
     target_dir = str(cur_dir.absolute()).split(target_directory)[0] + target_directory
     os.chdir(target_dir)
-    print(f'Current dir: {target_dir}')
+    print(f"Current dir: {target_dir}")
 
 
 def display_pretty_table_html(data: dict) -> None:
@@ -60,4 +60,4 @@ def print_table(data: Any, use_html: bool = True) -> None:
     if use_html:
         display(HTML(table))
     else:
-        print(table)
+        print(table)

speedy_utils/common/report_manager.py

@@ -3,7 +3,6 @@ from collections import defaultdict
 from datetime import datetime
 
 
-
 class ReportManager:
     def __init__(self):
         self.cache_dir = os.path.expanduser("~/.cache/speedy_utils")
@@ -41,7 +40,7 @@ class ReportManager:
             [
                 "\n### Results Overview",
                 f"- Total items processed: {len(results)}",
-                f"- Success rate: {(len(results) - len(errors))/len(results)*100:.1f}%",
+                f"- Success rate: {(len(results) - len(errors)) / len(results) * 100:.1f}%",
                 f"- Total errors: {len(errors)}",
             ]
         )
@@ -49,7 +48,7 @@ class ReportManager:
         if execution_time:
             md_content.append(f"- Execution time: {execution_time:.2f}s")
             md_content.append(
-                f"- Average speed: {len(results)/execution_time:.1f} items/second"
+                f"- Average speed: {len(results) / execution_time:.1f} items/second"
             )
 
         if error_groups:

speedy_utils/common/utils_cache.py

@@ -1,3 +1,4 @@
+import asyncio
 import functools
 import inspect
 import json
@@ -6,7 +7,7 @@ import os.path as osp
 import pickle
 import uuid
 from threading import Lock
-from typing import Any, Literal
+from typing import Any, Awaitable, Callable, Literal, TypeVar
 
 import cachetools
 import pandas as pd
@@ -26,6 +27,10 @@ thread_locker = Lock()
 disk_lock = Lock()
 mem_lock = Lock()
 
+# Add async-specific types
+T = TypeVar('T')
+AsyncFunc = Callable[..., Awaitable[T]]
+
 
 def compute_func_id(func, args, kwargs, ignore_self, keys):
     func_source = get_source(func)
@@ -144,6 +149,61 @@ def _disk_memoize(func, keys, cache_dir, ignore_self, verbose):
     return wrapper
 
 
+def _async_disk_memoize(func, keys, cache_dir, ignore_self, verbose):
+    @functools.wraps(func)
+    async def wrapper(*args, **kwargs):
+        try:
+            # Compute cache path as before
+            func_source, sub_dir, key_id = compute_func_id(
+                func, args, kwargs, ignore_self, keys
+            )
+            if func_source is None:
+                return await func(*args, **kwargs)
+            if sub_dir == "funcs":
+                cache_path = osp.join(cache_dir, sub_dir, func.__name__, key_id)
+            else:
+                cache_path = osp.join(cache_dir, sub_dir, key_id)
+            mkdir_or_exist(osp.dirname(cache_path))
+
+            # First check with disk lock (run in thread to avoid blocking)
+            def check_cache():
+                with disk_lock:
+                    if osp.exists(cache_path):
+                        try:
+                            return load_json_or_pickle(cache_path)
+                        except Exception as e:
+                            if osp.exists(cache_path):
+                                os.remove(cache_path)
+                            logger.opt(depth=1).warning(
+                                f"Error loading cache: {str(e)[:100]}, continue to recompute"
+                            )
+                    return None
+            
+            # Run cache check in thread pool to avoid blocking
+            loop = asyncio.get_event_loop()
+            cached_result = await loop.run_in_executor(None, check_cache)
+            if cached_result is not None:
+                return cached_result
+
+            result = await func(*args, **kwargs)
+
+            # Write result under disk lock (run in thread to avoid blocking)
+            def write_cache():
+                with disk_lock:
+                    if not osp.exists(cache_path):
+                        dump_json_or_pickle(result, cache_path)
+            
+            await loop.run_in_executor(None, write_cache)
+            return result
+        except Exception as e:
+            logger.opt(depth=1).warning(
+                f"Failed to cache {func.__name__}: {e}, continue to recompute without cache"
+            )
+            return await func(*args, **kwargs)
+
+    return wrapper
+
+
 def _memory_memoize(func, size, keys, ignore_self):
     global LRU_MEM_CACHE
     if LRU_MEM_CACHE.maxsize != size:
@@ -176,6 +236,38 @@ def _memory_memoize(func, size, keys, ignore_self):
     return wrapper
 
 
+def _async_memory_memoize(func, size, keys, ignore_self):
+    global LRU_MEM_CACHE
+    if LRU_MEM_CACHE.maxsize != size:
+        LRU_MEM_CACHE = cachetools.LRUCache(maxsize=size)
+
+    @functools.wraps(func)
+    async def wrapper(*args, **kwargs):
+        func_source, sub_dir, key_id = compute_func_id(
+            func, args, kwargs, ignore_self, keys
+        )
+        if func_source is None:
+            return await func(*args, **kwargs)
+        name = identify((func_source, sub_dir, key_id))
+
+        if not hasattr(func, "_mem_cache"):
+            func._mem_cache = LRU_MEM_CACHE
+
+        with mem_lock:
+            if name in func._mem_cache:
+                # logger.debug(f"Cache HIT (memory) for {func.__name__}, key={name}")
+                return func._mem_cache[name]
+
+        result = await func(*args, **kwargs)
+
+        with mem_lock:
+            if name not in func._mem_cache:
+                func._mem_cache[name] = result
+        return result
+
+    return wrapper
+
+
 def both_memoize(func, keys, cache_dir, ignore_self):
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
@@ -220,6 +312,63 @@ def both_memoize(func, keys, cache_dir, ignore_self):
     return wrapper
 
 
+def _async_both_memoize(func, keys, cache_dir, ignore_self):
+    @functools.wraps(func)
+    async def wrapper(*args, **kwargs):
+        func_source, sub_dir, key_id = compute_func_id(
+            func, args, kwargs, ignore_self, keys
+        )
+        if func_source is None:
+            return await func(*args, **kwargs)
+
+        mem_key = identify((func_source, sub_dir, key_id))
+        if not hasattr(func, "_mem_cache"):
+            func._mem_cache = LRU_MEM_CACHE
+
+        with mem_lock:
+            if mem_key in func._mem_cache:
+                # logger.debug(f"Cache HIT (memory) for {func.__name__}, key={mem_key}")
+                return func._mem_cache[mem_key]
+
+        if sub_dir == "funcs":
+            cache_path = osp.join(cache_dir, sub_dir, func.__name__, key_id)
+        else:
+            cache_path = osp.join(cache_dir, sub_dir, key_id)
+        mkdir_or_exist(osp.dirname(cache_path))
+
+        # Check disk cache in thread pool to avoid blocking
+        def check_disk_cache():
+            with disk_lock:
+                if osp.exists(cache_path):
+                    return load_json_or_pickle(cache_path)
+            return None
+
+        loop = asyncio.get_event_loop()
+        disk_result = await loop.run_in_executor(None, check_disk_cache)
+        
+        if disk_result is not None:
+            with mem_lock:
+                func._mem_cache[mem_key] = disk_result
+            return disk_result
+
+        # logger.debug(f"Cache MISS for {func.__name__}, key={cache_path}")
+        result = await func(*args, **kwargs)
+
+        # Write to disk in thread pool to avoid blocking
+        def write_disk_cache():
+            with disk_lock:
+                if not osp.exists(cache_path):
+                    dump_json_or_pickle(result, cache_path)
+
+        await loop.run_in_executor(None, write_disk_cache)
+        
+        with mem_lock:
+            func._mem_cache[mem_key] = result
+        return result
+
+    return wrapper
+
+
 def memoize(
     _func=None,
     *,
@@ -234,7 +383,17 @@ def memoize(
         cache_dir = osp.expanduser(cache_dir)
 
     def decorator(func):
+        # Check if function is async
+        is_async = inspect.iscoroutinefunction(func)
+        
         if cache_type == "memory":
+            if is_async:
+                return _async_memory_memoize(
+                    func,
+                    size,
+                    keys,
+                    ignore_self,
+                )
             return _memory_memoize(
                 func,
                 size,
@@ -242,6 +401,14 @@ def memoize(
                 ignore_self,
             )
         elif cache_type == "disk":
+            if is_async:
+                return _async_disk_memoize(
+                    func,
+                    keys,
+                    cache_dir,
+                    ignore_self,
+                    verbose,
+                )
             return _disk_memoize(
                 func,
                 keys,
@@ -249,6 +416,15 @@ def memoize(
                 ignore_self,
                 verbose,
             )
+        
+        # cache_type == "both"
+        if is_async:
+            return _async_both_memoize(
+                func,
+                keys,
+                cache_dir,
+                ignore_self,
+            )
         return both_memoize(
             func,
             keys,
@@ -256,9 +432,63 @@ def memoize(
             verbose,
         )
 
+    # Handle both @memoize and @memoize() usage patterns
     if _func is None:
         return decorator
-    return decorator(_func)
+    else:
+        return decorator(_func)
+
+
+def amemoize(
+    _func=None,
+    *,
+    keys: list[str] | None = None,
+    cache_dir: str = SPEED_CACHE_DIR,
+    cache_type: Literal["memory", "disk", "both"] = "disk",
+    size: int = 10240,
+    ignore_self: bool = True,
+    verbose: bool = False,
+):
+    """
+    Async-specific memoization decorator for coroutine functions.
+    
+    Args:
+        _func: The async function to memoize (when used without parentheses)
+        keys: Specific argument keys to use for cache key generation
+        cache_dir: Directory for disk cache storage
+        cache_type: Type of caching - "memory", "disk", or "both"  
+        size: Size of memory cache (for memory/both types)
+        ignore_self: Whether to ignore 'self' parameter in cache key
+        verbose: Enable verbose logging
+        
+    Returns:
+        Decorated async function with memoization
+        
+    Example:
+        @amemoize(cache_type="both")
+        async def my_async_func(x: int) -> str:
+            return str(x)
+    """
+    if "~/" in cache_dir:
+        cache_dir = osp.expanduser(cache_dir)
+
+    def decorator(func):
+        # Ensure the function is actually async
+        if not inspect.iscoroutinefunction(func):
+            raise ValueError(f"amemoize can only be used with async functions. {func.__name__} is not async.")
+        
+        if cache_type == "memory":
+            return _async_memory_memoize(func, size, keys, ignore_self)
+        elif cache_type == "disk":
+            return _async_disk_memoize(func, keys, cache_dir, ignore_self, verbose)
+        else:  # cache_type == "both"
+            return _async_both_memoize(func, keys, cache_dir, ignore_self)
+
+    # Handle both @amemoize and @amemoize() usage patterns
+    if _func is None:
+        return decorator
+    else:
+        return decorator(_func)
 
 
-__all__ = ["memoize", "identify", "identify_uuid"]
+__all__ = ["memoize", "identify", "identify_uuid", "amemoize"]

speedy_utils/common/utils_io.py

@@ -147,6 +147,8 @@ def jdumps(obj, ensure_ascii=False, indent=2, **kwargs):
     return json.dumps(obj, ensure_ascii=ensure_ascii, indent=indent, **kwargs)
 
 
+
+
 __all__ = [
     "dump_json_or_pickle",
     "dump_jsonl",

speedy_utils/scripts/mpython.py

@@ -91,9 +91,7 @@ def main():
         cpu_end = ((i + 1) * cpu_per_process - 1) % args.total_cpu
         ENV = f"CUDA_VISIBLE_DEVICES={gpu} MP_ID={i} MP_TOTAL={args.total_fold}"
         if taskset_path:
-            fold_cmd = (
-                f"{ENV} {taskset_path} -c {cpu_start}-{cpu_end}  {path_python} {cmd_str}"
-            )
+            fold_cmd = f"{ENV} {taskset_path} -c {cpu_start}-{cpu_end}  {path_python} {cmd_str}"
         else:
             fold_cmd = f"{ENV} {path_python} {cmd_str}"
 

{speedy_utils-1.1.5.dist-info → speedy_utils-1.1.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: speedy-utils
-Version: 1.1.5
+Version: 1.1.7
 Summary: Fast and easy-to-use package for data science
 Author: AnhVTH
 Author-email: anhvth.226@gmail.com

speedy_utils-1.1.7.dist-info/RECORD

@@ -0,0 +1,39 @@
+llm_utils/__init__.py,sha256=AYmJ297r0OjYmV1uNWFwznbqfuLTjCV7L2Ee12jxhpw,602
+llm_utils/chat_format/__init__.py,sha256=8dBIUqFJvkgQYedxBtcyxt-4tt8JxAKVap2JlTXmgaM,737
+llm_utils/chat_format/display.py,sha256=M-__JpcJSqjqeP4LiW7-yF8fVL37yUEUdaNC4VEgIo8,10181
+llm_utils/chat_format/transform.py,sha256=eU0c3PdAHCNLuGP1UqPwln0B34Lv3bt_uV9v9BrlCN4,5402
+llm_utils/chat_format/utils.py,sha256=xTxN4HrLHcRO2PfCTR43nH1M5zCa7v0kTTdzAcGkZg0,1229
+llm_utils/group_messages.py,sha256=Oe2tlhg-zRodG1-hodYebddrR77j9UdE05LzJw0EvYI,3622
+llm_utils/lm/__init__.py,sha256=rX36_MsnekM5GHwWS56XELbm4W5x2TDwnPERDTfo0eU,194
+llm_utils/lm/async_lm/__init__.py,sha256=PUBbCuf5u6-0GBUu-2PI6YAguzsyXj-LPkU6vccqT6E,121
+llm_utils/lm/async_lm/_utils.py,sha256=P1-pUDf_0pDmo8WTIi43t5ARlyGA1RIJfpAhz-gfA5g,6105
+llm_utils/lm/async_lm/async_llm_task.py,sha256=i5Z9yLBnxFK_937JHFMu7A-tImQnRuGAsnWDcC70syg,18648
+llm_utils/lm/async_lm/async_lm.py,sha256=J1KC7qCpG_CyJMWca4q71la7JHoANiLLSNQrQH44-z0,14045
+llm_utils/lm/async_lm/async_lm_base.py,sha256=BIIKAgmSZ9S5gLIf4AKGWg8N1eAr0l_fRfgG6UnKCK4,14839
+llm_utils/lm/async_lm/lm_specific.py,sha256=KmqdCm3SJ5MqN-dRJd6S5tq5-ve1X2eNWf2CMFtc_3s,3926
+llm_utils/lm/utils.py,sha256=a0KJj8vjT2fHKb7GKGNJjJHhKLThwpxIL7vnV9Fr3ZY,4584
+llm_utils/scripts/README.md,sha256=yuOLnLa2od2jp4wVy3rV0rESeiV3o8zol5MNMsZx0DY,999
+llm_utils/scripts/vllm_load_balancer.py,sha256=TT5Ypq7gUcl52gRFp--ORFFjzhfGlcaX2rkRv8NxlxU,37259
+llm_utils/scripts/vllm_serve.py,sha256=4NaqpVs7LBvxtvTCMPsNCAOfqiWkKRttxWMmWY7SitA,14729
+speedy_utils/__init__.py,sha256=ZtnitBT13OS3xjmsVoVHjmL5RIWaH12PMcp6UDHQjaE,5776
+speedy_utils/all.py,sha256=t-HKzDmhF1MTFnmq7xRnPs5nFG_aZaLH9Ua0RM6nQ9Y,4855
+speedy_utils/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+speedy_utils/common/clock.py,sha256=3n4FkCW0dz46O8By09V5Pve1DSMgpLDRbWEVRryryeQ,7423
+speedy_utils/common/function_decorator.py,sha256=BspJ0YuGL6elS7lWBAgELZ-sCfED_1N2P5fgH-fCRUQ,2132
+speedy_utils/common/logger.py,sha256=JqW9gG4ujfq4RldNeYP2p52BYgCwjkYeGGYyzLn6mfY,6422
+speedy_utils/common/notebook_utils.py,sha256=-97kehJ_Gg3TzDLubsLIYJcykqX1NXhbvBO6nniZSYM,2063
+speedy_utils/common/report_manager.py,sha256=eBiw5KY6bWUhwki3B4lK5o8bFsp7L5x28X9GCI-Sd1w,3899
+speedy_utils/common/utils_cache.py,sha256=G0M_iv3T8QqbBNNiS1LDz6MrRycQjiYLMzmHYpDUCjU,16348
+speedy_utils/common/utils_io.py,sha256=tfptex3pbmhXOftr__V-3DbhuDVSm01j4vg39R5jbwI,4792
+speedy_utils/common/utils_misc.py,sha256=cdEuBBpiB1xpuzj0UBDHDuTIerqsMIw37ENq6EXliOw,1795
+speedy_utils/common/utils_print.py,sha256=iQqnOYw2EFC8TqeSDbrcnIQAUKT7FbB8Mec8b2aGAzw,4833
+speedy_utils/multi_worker/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+speedy_utils/multi_worker/process.py,sha256=BI-sgzzQ0_N8kOfaS_3ZAGZ3d6panYzJ3-BGZthY4dQ,6824
+speedy_utils/multi_worker/thread.py,sha256=u_hTwXh7_FciMa5EukdEA1fDCY_vUC4moDceBXk2b6w,16326
+speedy_utils/scripts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+speedy_utils/scripts/mpython.py,sha256=IvywP7Y0_V6tWfMP-4MjPvN5_KfxWF21xaLJsCIayCk,3821
+speedy_utils/scripts/openapi_client_codegen.py,sha256=f2125S_q0PILgH5dyzoKRz7pIvNEjCkzpi4Q4pPFRZE,9683
+speedy_utils-1.1.7.dist-info/METADATA,sha256=EV1Jj7hpGp5Zf9hkj0chArF9iVmJUCUFeSFW2VkXzvI,7441
+speedy_utils-1.1.7.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+speedy_utils-1.1.7.dist-info/entry_points.txt,sha256=T1t85jwx8fK6m5msdkBGIXH5R5Kd0zSL0S6erXERPzg,237
+speedy_utils-1.1.7.dist-info/RECORD,,