lytools-HPC 0.0.1__py3-none-any.whl

lytools_HPC/__init__.py

@@ -0,0 +1,50 @@
+# coding='utf-8'
+__version__ = '0.0.1'
+__package_name__ = 'lytools_HPC'
+import warnings
+import requests
+from packaging import version
+
+def check_latest_version(package_name, current_version, cache_hours=0.1):
+    cache_file = os.path.expanduser(f"~/.{package_name}_version_check")
+    latest_version = None
+    now = time.time()
+
+    if os.path.exists(cache_file):
+        try:
+            with open(cache_file, "r") as f:
+                data = json.load(f)
+            if now - data["timestamp"] < cache_hours * 3600:
+                latest_version = data["latest_version"]
+        except Exception:
+            pass
+
+    if latest_version is None:
+        try:
+            url = f"https://pypi.org/pypi/{package_name}/json"
+            resp = requests.get(url, timeout=3)
+            if resp.status_code == 200:
+                data = resp.json()
+                latest_version = data["info"]["version"]
+                with open(cache_file, "w") as f:
+                    json.dump({"latest_version": latest_version, "timestamp": now}, f)
+        except Exception:
+            return
+
+    try:
+        if version.parse(current_version) < version.parse(latest_version):
+            warnings.warn(
+                f"{package_name} {current_version} is outdated. "
+                f"Latest version is {latest_version}. "
+                f"Upgrade via: pip install -U {package_name}"
+            )
+    except Exception:
+        pass
+
+try:
+    current = __version__
+    check_latest_version(__package_name__, current)
+except Exception:
+    pass
+
+from ._lytools_HPC import *

lytools_HPC/_lytools_HPC.py

@@ -0,0 +1,450 @@
+import submitit
+from concurrent.futures import ThreadPoolExecutor
+import redis
+from threading import Lock
+from pathos.multiprocessing import ProcessPool
+import time
+from rich.text import Text
+
+import os
+import json
+from pathlib import Path
+import pickle
+
+from rich.progress import (
+    Progress,
+    BarColumn,
+    TextColumn,
+    TimeRemainingColumn,
+    TimeElapsedColumn,
+    ProgressColumn
+)
+
+
+def listdir(fdir):
+    '''
+    Mac OS
+    list the names of the files in the directory
+    return sorted files list without '.DS_store'
+    '''
+    list_dir = []
+    for f in sorted(os.listdir(fdir)):
+        if f.startswith('.'):
+            continue
+        list_dir.append(f)
+    return list_dir
+
+
+def mkdir(dir, force=False):
+    if not os.path.isdir(dir):
+        if force == True:
+            os.makedirs(dir)
+        else:
+            os.mkdir(dir)
+
+def split_into_n_jobs(lst, n_jobs):
+    """
+    Split the list into n_jobs parts.
+    """
+    total = len(lst)
+    chunk_size = total // n_jobs
+    remainder = total % n_jobs
+
+    chunks = []
+    start = 0
+
+    for i in range(n_jobs):
+        extra = 1 if i < remainder else 0
+        end = start + chunk_size + extra
+        chunks.append(lst[start:end])
+        start = end
+    return chunks
+
+_REDIS = None
+
+def get_redis():
+    global _REDIS
+    if _REDIS is None:
+        _REDIS = HPC_redis().conn_redis()
+    return _REDIS
+
+_LOCK = Lock()
+_LOCAL_COUNTER = 0
+
+def update_i(job_name, batch_size=100):
+    '''
+    :param job_name:
+    '''
+    global _LOCAL_COUNTER
+
+    with _LOCK:
+        _LOCAL_COUNTER += 1
+
+        if _LOCAL_COUNTER >= batch_size:
+            r = get_redis()
+            r.hincrby(job_name, "step", _LOCAL_COUNTER)
+            _LOCAL_COUNTER = 0
+
+class IterSpeedColumn(ProgressColumn):
+
+    def render(self, task):
+        if task.finished:
+            return Text(f"{task.speed:.2f} it/s", style="green")
+
+        speed = task.speed
+
+        if speed is None:
+            return Text("-- it/s", style="dim")
+
+        return Text(f"{speed:.2f} it/s", style="cyan")
+
+def sumbit_jobs_array(func,params_list,log_folder,job_name,
+                        job_number_limit=500,
+                        parallel_process_per_task=10,
+                        slurm_array_parallelism=20,
+                        parallel_process_p_or_t='p',
+                        cpus_per_task=1,
+                        mem_gb=1,
+                        timeout_min=5,
+                        slurm_partition="general",
+                        exclude_nodes=None,
+                        pbar_update_freq=1
+                      ):
+    '''
+    :param func: the kernel function to run, should take one argument, e.g. func(params)
+    :param params_list: list of tuples [params1, params2, ...]
+    :param log_folder: slurm log_folder
+    :param job_name: slurm job_name
+    :param job_number_limit: number of total jobs you want to submit to slurm
+    :param parallel_process_per_task: number of parallel processes per task, Recommend equal to :param cpus_per_task
+    :param slurm_array_parallelism: slurm array parallelism
+    :param parallel_process_p_or_t: 'p' for multiprocessing, 't' for multi-threading
+    :param cpus_per_task: number of cpus per task
+    :param mem_gb: memory per task
+    :param timeout_min: timeout in minutes
+    :param slurm_partition: slurm partition
+    :param exclude_nodes: list of nodes to exclude
+    :param pbar_update_freq: frequency of updating progress bar
+    '''
+
+    if len(params_list) == 0:
+        raise ValueError("params_list is empty")
+    if len(params_list) > job_number_limit:
+        super_params_list = split_into_n_jobs(params_list,job_number_limit)
+        if parallel_process_p_or_t == 't':
+            def super_func(chunk):
+                def wrapper(p):
+                    func(p)
+                    update_i(job_name,pbar_update_freq)
+
+                with ThreadPoolExecutor(max_workers=parallel_process_per_task) as Thread_:
+                    list(Thread_.map(wrapper, chunk))
+
+        elif parallel_process_p_or_t == 'p':
+            def super_func(chunk):
+                # for p in chunk:
+                #     func(p)
+
+                def wrapper(p):
+                    func(p)
+                    update_i(job_name,pbar_update_freq)
+
+                pool = ProcessPool(nodes=parallel_process_per_task)
+                pool.map(wrapper, chunk)
+            pass
+        else:
+            raise ValueError("parallel_process_p_or_t must be 'p' for multiprocessing or 't' for threading")
+
+        final_params_list = super_params_list
+        final_func = super_func
+    else:
+        final_params_list = params_list
+        final_func = func
+
+    if os.path.exists(log_folder):
+        for f in os.listdir(log_folder):
+            os.remove(os.path.join(log_folder, f))
+    mkdir(log_folder, force=True)
+
+    print('submiting...')
+    executor = submitit.AutoExecutor(folder=log_folder)
+    executor.update_parameters(
+        slurm_job_name=job_name,
+        cpus_per_task=cpus_per_task,
+        mem_gb=mem_gb,
+        timeout_min=timeout_min,
+        slurm_array_parallelism=slurm_array_parallelism,
+        slurm_partition=slurm_partition,
+        exclude=exclude_nodes,
+    )
+    jobs = executor.map_array(final_func, final_params_list)
+
+    info = {
+        "Total Cores Used": cpus_per_task * slurm_array_parallelism,
+        "Concurrent Processes": parallel_process_per_task * slurm_array_parallelism,
+        "Total Loop Length": len(params_list),
+        "Number of Jobs": len(final_params_list),
+        "First Job ID": jobs[0].job_id,
+        "Memory for Each Job(GB)":mem_gb,
+        "Time out Minutes For Each Job": timeout_min,
+        "Partition":slurm_partition,
+    }
+    pretty_table_print(info)
+
+
+
+def pretty_table_print(info):
+    max_key_len = max(len(k) for k in info)
+
+    print("\n=== Job Summary ===")
+    for k, v in info.items():
+        print(f"{k + ':':<{max_key_len + 2}} {v}")
+    print("=" * (max_key_len + 15))
+    pass
+
+def monitoring_job(progress_dir_json):
+    refresh_interval = 1
+    # refresh_interval = 5
+    task_map = {}
+    with Progress(
+            TextColumn("[bold yellow]{task.description}"),
+            BarColumn(style="red",  # Unfinished part
+                      complete_style="green",  # Finished part
+                      finished_style="yellow"),  # 100% finished),
+            TextColumn("{task.completed}/{task.total} {task.percentage:>3.0f}%"),
+            TimeElapsedColumn(),
+            TimeRemainingColumn(),
+    ) as progress_total:
+
+        with Progress(
+                TextColumn("[bold yellow]{task.description}"),
+                BarColumn(),
+                TextColumn("{task.completed}/{task.total} {task.percentage:>3.0f}%"),
+                TimeElapsedColumn(),
+                TimeRemainingColumn(),
+        ) as progress:
+
+            overall_task = None
+            total_jobs = None
+
+            while True:
+
+                files = list(progress_dir_json.glob("*.json"))
+                finished_jobs = 0
+
+                for f in files:
+
+                    try:
+                        with open(f) as fp:
+                            data = json.load(fp)
+                    except:
+                        continue
+
+                    jobid = data["jobid"]
+                    step = data["step"]
+                    total = data["total"]
+                    total_job = data["total_job"]
+                    sub_job_name = data["sub_job_name"]
+
+                    # 初始化 total progress
+
+                    if total_jobs is None:
+                        total_jobs = total_job
+                        overall_task = progress_total.add_task(
+                            "[bold green]TOTAL", total=total_jobs
+                        )
+
+                    # if completed
+                    if step >= total:
+
+                        finished_jobs += 1
+
+                        if jobid in task_map:
+                            progress.remove_task(task_map[jobid])
+                            del task_map[jobid]
+
+                        continue
+
+                    # create bar
+                    if jobid not in task_map:
+                        task_map[jobid] = progress.add_task(
+                            f"({jobid}/{total_job}) {sub_job_name}", total=total
+                        )
+
+                    progress.update(task_map[jobid], completed=step)
+
+                # update total bar
+                if overall_task is not None:
+                    progress_total.update(overall_task, completed=finished_jobs)
+
+                # All tasks completed
+                if total_jobs and finished_jobs == total_jobs:
+                    break
+
+                time.sleep(refresh_interval)
+    pass
+
+class HPC_redis:
+
+    def __init__(self):
+        self.r = self.conn_redis()
+        pass
+
+    def conn_redis(self):
+        redis_conf = Path.home() / '.config' / 'redis' / 'redis.conf'
+        with open(redis_conf) as f:
+            redis_conf = f.readlines()
+            host = redis_conf[0].strip()
+            port = int(redis_conf[1].strip())
+            passwd = redis_conf[2].strip()
+
+        r = redis.Redis(
+            host=host,
+            port=port,
+            password=passwd,
+        )
+        # print(f"Connected to Redis at {host}:{port}")
+        return r
+
+    def hit_redis(self,job_name,amount=1):
+        self.r.hincrby(name=job_name,key=job_name,amount=amount)
+
+    def set_total_num(self,job_name,total_job:int):
+        # r.delete(job_name+'_total')
+        # r.hincrby(name=job_name, key=task_name+'_total', amount=total_job)
+        self.r.hset(job_name, 'total', str(total_job))
+
+    def delete_job(self,job_name):
+        self.r.delete(job_name)
+        pass
+
+    def query_redis(self,job_name):
+        while True:
+            print(self.r.hgetall(job_name))
+            time.sleep(1)
+
+class Check_logs:
+    def __init__(self,log_folder):
+        self.log_folder = log_folder
+        pass
+
+    def read_err_files(self):
+        log_folder = self.log_folder
+        log_folder = Path(log_folder)
+        err_count = 0
+        for f in listdir(log_folder):
+            if not f.endswith(".err"):
+                continue
+            fpath = log_folder / f
+            # print(fpath)
+            with open(fpath) as fr:
+                err_content = fr.read()
+                if len(err_content) != 0:
+                    print('==============')
+                    print(f"Error in file: {f}")
+                    print(err_content)
+                    err_count += 1
+        print('#################')
+        print(f'Total error logs: {err_count}')
+        print('#################')
+        pass
+
+    def read_out_files(self):
+        log_folder = self.log_folder
+        log_folder = Path(log_folder)
+        count = 0
+        for f in listdir(log_folder):
+            if not f.endswith(".out"):
+                continue
+            fpath = log_folder / f
+            # print(fpath)
+            with open(fpath) as fr:
+                log_content = fr.read()
+                print(log_content)
+                print('------------')
+                count += 1
+        print(f'Total files: {count}')
+        pass
+
+    def read_result_files(self):
+        log_folder = self.log_folder
+        log_folder = Path(log_folder)
+        count = 0
+        for f in listdir(log_folder):
+            if not f.endswith("_result.pkl"):
+                continue
+            fpath = log_folder / f
+            # print(fpath)
+            content = pickle.load(open(fpath, 'rb'))
+            print(content)
+            print('------------')
+            count += 1
+        print(f'Total files: {count}')
+        pass
+
+    def read_submit_files(self):
+        log_folder = self.log_folder
+        log_folder = Path(log_folder)
+        count = 0
+        for f in listdir(log_folder):
+            if not f.endswith("_submitted.pkl"):
+                continue
+            fpath = log_folder / f
+            # print(fpath)
+            content = pickle.load(open(fpath, 'rb'))
+            print(content)
+            print('------------')
+            count += 1
+        print(f'Total files: {count}')
+        pass
+
+
+def init_job(job_name,param_list):
+    total_job = len(param_list)
+    r = HPC_redis().conn_redis()
+    r.delete(job_name)
+    r.hset(job_name, 'total', str(total_job))
+    r.hset(job_name, 'step', str(0))
+
+
+def flush_progress(job_name):
+    global _LOCAL_COUNTER
+
+    if _LOCAL_COUNTER > 0:
+        r = get_redis()
+        r.hincrby(job_name, "step", _LOCAL_COUNTER)
+        _LOCAL_COUNTER = 0
+
+
+def update_progress(job_name, batch_size=100):
+    update_i(job_name, batch_size)
+    flush_progress(job_name)
+
+def progress_bar_monitoring(job_name):
+    hpc_redis = HPC_redis()
+    info=hpc_redis.r.hgetall(job_name)
+    total = info.get(b'total')
+    step = info.get(b'step')
+    total_int = int(total)
+    step_int = int(step)
+    # exit()
+    with Progress(
+        TextColumn("[bold yellow]{task.description}"),
+        BarColumn(),
+        TextColumn("{task.completed}/{task.total} {task.percentage:>3.0f}%"),
+        TimeElapsedColumn(),
+        IterSpeedColumn(),
+        TimeRemainingColumn(),
+        # show_speed=True,
+    ) as progress:
+        task = progress.add_task(f"[bold green]{job_name}", total=total_int)
+        while True:
+            info = hpc_redis.r.hgetall(job_name)
+            # print(info)
+            step = info.get(b'step')
+            step_int = int(step)
+            progress.update(task, completed=step_int)
+            if step_int >= total_int:
+                break
+            time.sleep(1)
+

lytools_hpc-0.0.1.dist-info/METADATA

@@ -0,0 +1,119 @@
+Metadata-Version: 2.4
+Name: lytools_HPC
+Version: 0.0.1
+Home-page: https://github.com/leeyang1991/lytools_HPC
+Author: Yang Li
+Author-email: leeyang1991@gmail.com
+Requires-Python: >=3
+Description-Content-Type: text/markdown
+Requires-Dist: submitit
+Requires-Dist: redis
+Requires-Dist: pathos
+Requires-Dist: rich
+Requires-Dist: requests
+Requires-Dist: packaging
+Dynamic: author
+Dynamic: author-email
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: requires-dist
+Dynamic: requires-python
+
+#  HPC Functions
+
+## Monitoring you jobs with a progress bar. 
+
+This requires a Redis server to store the progress information. The following steps will guide you through setting up the Redis server and using submitit with a progress bar.
+### Python packages required:
+```bash
+pip install submitit redis pathos rich 
+```
+### Step 1: Deploy a Redis service
+**DO NOT** deploy it on a HPC nodes. You can use your local machine or a server with a public IP address.
+#### Deploy redis via Docker
+Compose file for Redis server:
+```yaml
+services:
+  redis-progress:
+    image: redis:7
+    container_name: redis-progress
+    restart: unless-stopped
+
+    ports:
+      - "6379:6379"
+
+    volumes:
+      - ./redis_data:/data
+
+    command: >
+      redis-server
+      --appendonly yes
+      --requirepass yourpassword
+      --maxmemory 2gb
+      --maxmemory-policy allkeys-lru
+      
+    deploy:
+      resources:
+        limits:
+          memory: 3g
+
+```
+```docker compose up -d``` # Start the Redis server
+
+#### Test Redis connection
+```python
+import redis
+r = redis.Redis(host='your_redis_host', port=6379, password='yourpassword')
+r.set('test_key', 'test_value')
+print(r.get('test_key').decode())
+```
+
+#### Create redis connection configuration file
+Location: ```~/.config/redis/redis.conf```
+```txt
+your_redis_host
+6379
+yourpassword
+```
+
+### Step 2: Submit your jobs
+```python
+from HPC_func import *
+
+def my_function(x):
+    return x * x
+
+job_name = 'Your_job_name'
+params_list = [1,2,3,4,5,6,7,8,9,10]
+log_folder = 'Your_log_folder'
+init_job(job_name, params_list)
+sumbit_jobs_array(my_function, params_list, log_folder,
+                  job_name=job_name,
+                  job_number_limit=1,
+                  parallel_process_per_task=2,
+                  slurm_array_parallelism=1,
+                  parallel_process_p_or_t='p',
+                  cpus_per_task=10,
+                  mem_gb=2,
+                  timeout_min=10,
+                  slurm_partition="general",
+                  exclude_nodes=None,
+                  pbar_update_freq=1,
+                  )
+```
+
+### Step 3: Monitor the progress bar
+```python
+from HPC_func import *
+job_name = 'Your_job_name'
+progress_bar_monitoring(job_name)
+```
+
+### Step 4: Check the log files
+```python
+from HPC_func import *
+log_folder = 'Your_log_folder'
+Check_logs(log_folder).read_err_files() # Read error files
+Check_logs(log_folder).read_out_files() # Read output files
+```

lytools_hpc-0.0.1.dist-info/RECORD

@@ -0,0 +1,6 @@
+lytools_HPC/__init__.py,sha256=4EY45Ok2uAxesY3AQo-puVc7rpfv6-dh7mA74GOatfo,1573
+lytools_HPC/_lytools_HPC.py,sha256=RirBkJmKsCoocgBnY68Ni5a5WtOqtmOG9LxrFOHe9Qg,13578
+lytools_hpc-0.0.1.dist-info/METADATA,sha256=QrnBDjy4_uCMgCkuN6ZYv5Bqc_pxTkPqohndO23iP2g,2944
+lytools_hpc-0.0.1.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+lytools_hpc-0.0.1.dist-info/top_level.txt,sha256=VIE9YTZbofWmUjTl3AFFJGfmWTOX7EV68IYJMMwsQuk,12
+lytools_hpc-0.0.1.dist-info/RECORD,,

lytools_hpc-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

lytools_hpc-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+lytools_HPC