taskcore-lib 0.0.2__py3-none-any.whl

taskcore/__init__.py

@@ -0,0 +1 @@
+from .base import FileSystemTaskQueueClient

taskcore/base.py

@@ -0,0 +1,233 @@
+import json
+import uuid
+from pathlib import Path
+import os
+import time
+import os
+import time
+
+import os
+import json
+import uuid
+import time
+from pathlib import Path
+from typing import Optional, Dict
+from taskcore.dist import atomic_write, get_unique_job_id,  wait_for_tid_file, get_tid_file_path
+import traceback
+from typing import Callable
+
+class FileSystemTaskQueue:
+    def __init__(self, base_dir: str):
+        self.base_dir = Path(base_dir)
+        self.pending_dir = self.base_dir / "pending"
+        self.running_dir = self.base_dir / "running"
+        self.finished_dir = self.base_dir / "finished"
+        # Ensure directories exist
+        for d in [self.pending_dir, self.running_dir, self.finished_dir]:
+            d.mkdir(parents=True, exist_ok=True)
+
+    def add_task(self, task_dict: Dict):
+        task_id = str(uuid.uuid4())
+        task_file = self.pending_dir / f"{task_id}.json"
+        with open(task_file, "w") as f:
+            json.dump(task_dict, f)
+        return task_id
+
+    def num_pending_tasks(self) -> int:
+        return len(list(self.pending_dir.iterdir()))
+    
+    def num_running_tasks(self) -> int:
+        return len(list(self.running_dir.iterdir()))
+    
+    def num_finished_tasks(self) -> int:
+        return len(list(self.finished_dir.iterdir()))
+    
+    def read_task(self, task_id: str) -> Dict:
+        task_file = self.running_dir / f"{task_id}.json"
+        with open(task_file, "r") as f:
+            return json.load(f)
+    
+    def fetch_task(self) -> Optional[str]:
+        for task_file in self.pending_dir.iterdir():
+            running_file = self.running_dir / task_file.name
+            try:
+                os.rename(task_file, running_file)  # atomic move
+                return str(running_file)
+            except FileNotFoundError:
+                continue  # another worker got it
+        return None
+
+    def finish_task(self, running_file: str):
+        running_path = Path(running_file)
+        finished_file = self.finished_dir / running_path.name
+        os.rename(running_path, finished_file)
+
+    def reset_task_timer(self, running_file: str):
+        now = time.time()
+        os.utime(running_file, (now, now))
+
+    def reclaim_stale_tasks(self, timeout_seconds: int = 600):
+        now = time.time()
+        for task_file in self.running_dir.iterdir():
+            if now - task_file.stat().st_mtime > timeout_seconds:
+                pending_file = self.pending_dir / task_file.name
+                try:
+                    os.rename(task_file, pending_file)
+                except FileNotFoundError:
+                    continue  # Already reclaimed
+
+    def get_task_dict(self, task_file: str) -> Dict:
+        with open(task_file, "r") as f:
+            return json.load(f)
+
+    def update_task_info(
+        self,
+        task_file: str,
+        new_dict: Optional[Dict] = None,
+    ):
+        """
+        Overwrite the entire task file with new_dict if provided.
+        Does nothing if new_dict is None.
+        """
+        if new_dict is not None:
+            path = Path(task_file)
+            with open(path, "w") as f:
+                json.dump(new_dict, f)
+        
+    
+class TaskQueueClient:
+    def __init__(self, queue: FileSystemTaskQueue, timeout: int = 60*60*4):
+        self.queue = queue
+        self.current_task_file = None
+        self.timeout = 60*60*4
+
+    def num_pending_tasks(self) -> int:
+        return self.queue.num_pending_tasks()
+    
+    def num_running_tasks(self) -> int:
+        return self.queue.num_running_tasks()
+    
+    def num_finished_tasks(self) -> int:    
+        return self.queue.num_finished_tasks()
+    
+    def read_task(self, task_file: str):
+ 
+        return self.queue.read_task(task_file)
+    
+    def fetch_task(self):
+        if self.current_task_file is not None:
+            raise RuntimeError("A task is already fetched. Finish or release it before fetching another.")
+        self.queue.reclaim_stale_tasks(self.timeout)  # Reclaim stale tasks before fetching
+        task_file = self.queue.fetch_task()
+        if task_file:
+            self.current_task_file = task_file
+            return self.current_task_file, self.get_current_task()
+        return None
+
+    def edit_current_task(self, new_dict):
+        if self.current_task_file is None:
+            raise RuntimeError("No task is currently fetched.")
+        self.queue.update_task_info(self.current_task_file, new_dict)
+
+    def finish_current_task(self):
+        if self.current_task_file is None:
+            raise RuntimeError("No task is currently fetched.")
+        self.queue.finish_task(self.current_task_file)
+        self.current_task_file = None
+
+    def release_current_task(self):
+        if self.current_task_file is None:
+            raise RuntimeError("No task is currently fetched.")
+        # Move the task back to pending
+        task_path = Path(self.current_task_file)
+        pending_file = self.queue.pending_dir / task_path.name
+        os.rename(task_path, pending_file)
+        self.current_task_file = None
+
+    def get_current_task(self):
+        if self.current_task_file is None:
+            return None
+        return self.queue.get_task_dict(self.current_task_file)
+
+    def has_task(self):
+        return self.current_task_file is not None
+        
+        
+def dummy_run_func(*args,**kwargs):
+    pass
+
+class FileSystemTaskQueueClient(TaskQueueClient):
+    def __init__(self,base_dir: str, rank: int = 0, timeout: int = 60*60*4):
+        self.queue = FileSystemTaskQueue(base_dir)
+        super().__init__(self.queue, timeout)
+        self.rank = rank
+
+    def add_task(self, task_dict: Dict):
+        return self.queue.add_task(task_dict)
+    
+    def fetch_task(self):
+        job_id = get_unique_job_id()
+        
+        if self.rank == 0:
+            # Main process fetches the next task
+            result =super().fetch_task() ## fetch the task on main process
+            if not result:
+                print("No task found")
+                raise RuntimeError("No task found")
+            
+            print(f"Fetched task {result}")
+            tid, config = result
+
+            shared_spot_name = get_tid_file_path(tid) ## get path for the shared spot
+            atomic_write(tid, shared_spot_name)## write to a shared spot to tell other processes to fetch the task
+
+            return tid, config
+        else:
+            time.sleep(15)
+            tid = wait_for_tid_file(job_id) ## wait for new file in shared spot.
+            
+            tid = tid.split("_")[-1].split(".")[0] ## get the tid from the file name
+
+            config = self.read_task(tid) ## broadcast the task to all processes
+            
+            return tid, config
+        
+    def fetch_and_run_task(self, init_func: Callable, func: Callable = dummy_run_func):
+        
+        tid, config = self.fetch_task() ## takes care of distributing a single task to all processes
+     
+        print(f"Running task {config}")
+
+        try:
+            trainer, extra_info = init_func(**config)
+            
+            if self.rank == 0:
+                config.update(extra_info)
+                #config["wandb_link"] = trainer.wandb_link
+                config["task_id"] = tid
+                self.edit_current_task(config)
+            
+            func(
+                trainer,
+            )
+            
+            if self.rank==0:
+                self.finish_current_task()
+                
+            return True
+                
+        except Exception as e:
+            print(f"Error running task {config}: {e}")
+            #print("traceback:", e.__traceback__)
+            traceback.print_exc()
+            if self.rank==0:
+                self.release_current_task()
+                
+            return False
+        finally:
+            
+            shared_spot_name = get_tid_file_path(tid)
+            # Clean up the tid file after use
+            if self.rank==0 and shared_spot_name and os.path.exists(shared_spot_name):
+                os.remove(shared_spot_name)
+                

taskcore/cli.py

@@ -0,0 +1,39 @@
+from taskcore import FileSystemTaskQueueClient
+import json
+from termcolor import colored
+
+
+def main(base_dir: str, mode:str="status"):
+    client = FileSystemTaskQueueClient(base_dir)
+    
+    if mode == "status":
+        print(colored(f"Pending tasks: {client.num_pending_tasks()}", 'red', attrs=['bold']))
+        print(colored(f"Running tasks: {client.num_running_tasks()}", 'green', attrs=['bold']))
+        print(colored(f"Finished tasks: {client.num_finished_tasks()}", 'blue', attrs=['bold']))
+        
+    elif mode == "show":
+        
+        print(colored(f"Pending tasks: {client.num_pending_tasks()}", 'red', attrs=['bold']))
+        for task_file in client.queue.pending_dir.iterdir():
+            #print("--"*40)
+            task_dict = client.queue.get_task_dict(task_file)
+            print(json.dumps(task_dict, indent=4))
+            print(f"[{colored(task_file.name, 'green')}] {json.dumps(task_dict, indent=4)}")
+            
+        # colored(e.get('variant'),'green')
+        print(colored(f"Running tasks: {client.num_running_tasks()}", 'green', attrs=['bold']))
+        for task_file in client.queue.running_dir.iterdir():
+            #print("--"*40)
+            
+            task_dict = client.queue.get_task_dict(task_file)
+            print(f"[{colored(task_file.name, 'green')}]{json.dumps(task_dict, indent=4)}")
+            #print(json.dumps(task_dict, indent=4))
+            
+    else:
+        raise ValueError(f"Invalid mode: {mode}")
+ 
+
+
+if __name__ == "__main__":
+    import fire
+    fire.Fire(main)

taskcore/dist.py

@@ -0,0 +1,53 @@
+
+
+import os
+import glob
+import time
+
+def atomic_write(data, filename):
+    tmp_filename = filename + ".tmp"
+    with open(tmp_filename, "w") as f:
+        f.write(data)
+    os.rename(tmp_filename, filename)  # Atomic on POSIX
+
+def atomic_read(filename):
+    with open(filename, "r") as f:
+        return f.read().strip()
+    
+
+def wait_for_tid_file(job_id, timeout=60):
+    pattern = f"/tmp/{job_id}_*.json"
+    start = time.time()
+    tid_file = None
+    while True:
+        files = [f for f in os.listdir("/tmp") if f.startswith(f"{job_id}_") and f.endswith(".json")]
+        if files:
+            # Get the youngest file (most recently modified)
+            file_paths = [os.path.join("/tmp", f) for f in files]
+            tid_file = max(file_paths, key=os.path.getmtime)
+            # Wait for file to be non-empty
+            if os.path.getsize(tid_file) > 0:
+                return tid_file
+        if time.time() - start > timeout:
+            raise TimeoutError("Timeout waiting for tid file")
+        time.sleep(0.1)
+        
+def get_tid_file_path(tid):
+    tid_filename = os.path.basename(tid)
+    unique_id = os.environ.get("SLURM_JOB_ID") or os.environ.get("MY_JOB_ID")
+    if unique_id is None:
+        raise RuntimeError("No unique job ID found in environment!")
+    return f"/tmp/{unique_id}_{tid_filename}"
+
+import uuid
+import time
+
+def get_unique_job_id():
+    # Use SLURM_JOB_ID if available, else MY_JOB_ID, else generate and set MY_JOB_ID
+    job_id = os.environ.get("SLURM_JOB_ID")
+    if not job_id:
+        job_id = os.environ.get("MY_JOB_ID")
+        if not job_id:
+            job_id = str(uuid.uuid4())
+            os.environ["MY_JOB_ID"] = job_id
+    return job_id

taskcore_lib-0.0.2.dist-info/METADATA

@@ -0,0 +1,244 @@
+Metadata-Version: 2.4
+Name: taskcore-lib
+Version: 0.0.2
+Summary: Package for basic task queue in the filesystem.
+Home-page: https://github.com/goncalorafaria/taskcore
+Author: Goncalo Faria
+Author-email: gfaria@cs.washington.edu
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.6.0
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license-file
+Dynamic: requires-python
+Dynamic: summary
+
+# TaskCore
+
+A lightweight, filesystem-based task queue system for Python. TaskCore provides a simple way to distribute and manage tasks across multiple processes or machines using the filesystem as the backend.
+
+## Features
+
+- **Simple Setup**: No external dependencies - uses only Python standard library
+- **Filesystem-based**: Tasks are stored as JSON files on disk
+- **Multi-process Support**: Built-in support for distributed task processing
+- **Fault Tolerant**: Automatic task recovery and stale task reclamation
+- **Atomic Operations**: Safe concurrent access using atomic file operations
+- **SLURM Integration**: Works seamlessly with SLURM job schedulers
+
+## Installation
+
+```bash
+pip install taskcore
+```
+
+Or install from source:
+
+```bash
+git clone https://github.com/goncalorafaria/taskcore
+cd taskcore
+pip install -e .
+```
+
+## Quick Start
+
+### Basic Usage
+
+```python
+from taskcore import FileSystemTaskQueueClient
+
+# Initialize the task queue
+queue = FileSystemTaskQueueClient("/path/to/task/directory")
+
+# Add tasks to the queue
+task_config = {
+    "learning_rate": 1e-4,
+    "batch_size": 32,
+    "epochs": 100,
+    "model_name": "bert-base-uncased"
+}
+
+task_id = queue.add_task(task_config)
+print(f"Added task with ID: {task_id}")
+
+# Process tasks
+def init_function(**config):
+    # Initialize your model/trainer here
+    trainer = create_trainer(**config)
+    return trainer, {"status": "initialized"}
+
+def run_function(trainer):
+    # Run your training/processing here
+    trainer.train()
+
+# Fetch and run tasks
+while queue.fetch_and_run_task(init_func=init_function, func=run_function):
+    pass
+```
+
+### Multi-Process Example
+
+```python
+import os
+from taskcore import FileSystemTaskQueueClient
+
+def main():
+    # Use LOCAL_RANK for multi-process setups (e.g., with torch.distributed)
+    rank = int(os.environ.get("LOCAL_RANK", 0))
+    queue = FileSystemTaskQueueClient("/path/to/tasks", rank=rank)
+    
+    def init_func(**config):
+        # Initialize your distributed training setup
+        trainer = create_distributed_trainer(**config)
+        return trainer, {"wandb_link": trainer.wandb_link}
+    
+    def run_func(trainer):
+        trainer.train()
+    
+    # Process tasks until none are available
+    while queue.fetch_and_run_task(init_func=init_func, func=run_func):
+        pass
+
+if __name__ == "__main__":
+    main()
+```
+
+### Task Generation Example
+
+```python
+from taskcore import FileSystemTaskQueueClient
+
+# Initialize queue
+queue = FileSystemTaskQueueClient("/path/to/task/directory")
+
+# Define hyperparameter sweep
+learning_rates = [1e-4, 1e-5, 1e-6]
+batch_sizes = [16, 32, 64]
+model_names = ["bert-base-uncased", "roberta-base"]
+
+# Generate all combinations
+for lr in learning_rates:
+    for batch_size in batch_sizes:
+        for model_name in model_names:
+            config = {
+                "learning_rate": lr,
+                "batch_size": batch_size,
+                "model_name": model_name,
+                "epochs": 100,
+                "wandb_project": "my-experiment"
+            }
+            
+            task_id = queue.add_task(config)
+            print(f"Added task {task_id}: {config}")
+```
+
+## API Reference
+
+### FileSystemTaskQueueClient
+
+The main client class for interacting with the task queue.
+
+#### Constructor
+
+```python
+FileSystemTaskQueueClient(base_dir: str, rank: int = 0)
+```
+
+- `base_dir`: Directory where task files are stored
+- `rank`: Process rank for multi-process setups (default: 0)
+
+#### Methods
+
+##### `add_task(task_dict: Dict) -> str`
+Add a new task to the queue.
+
+```python
+task_id = queue.add_task({"param": "value"})
+```
+
+##### `fetch_and_run_task(init_func: Callable, func: Callable) -> bool`
+Fetch a task, initialize it with `init_func`, and run it with `func`.
+
+```python
+def init_func(**config):
+    return trainer, extra_info
+
+def run_func(trainer):
+    trainer.train()
+
+success = queue.fetch_and_run_task(init_func, run_func)
+```
+
+##### `fetch_task() -> Tuple[str, Dict]`
+Fetch the next available task.
+
+```python
+task_id, config = queue.fetch_task()
+```
+
+##### `finish_current_task()`
+Mark the current task as completed.
+
+##### `release_current_task()`
+Release the current task back to the pending queue.
+
+##### `edit_current_task(new_dict: Dict)`
+Update the current task's configuration.
+
+##### `get_current_task() -> Dict`
+Get the configuration of the current task.
+
+##### `has_task() -> bool`
+Check if a task is currently being processed.
+
+## Directory Structure
+
+TaskCore creates the following directory structure:
+
+```
+base_dir/
+├── pending/     # Tasks waiting to be processed
+├── running/     # Tasks currently being processed
+└── finished/    # Completed tasks
+```
+
+## Environment Variables
+
+TaskCore uses these environment variables for distributed processing:
+
+- `SLURM_JOB_ID`: Job ID from SLURM scheduler
+- `MY_JOB_ID`: Fallback job ID if SLURM_JOB_ID is not available
+- `LOCAL_RANK`: Process rank in multi-process setups
+
+## Error Handling
+
+Tasks that fail during execution are automatically released back to the pending queue:
+
+```python
+try:
+    queue.fetch_and_run_task(init_func, run_func)
+except Exception as e:
+    print(f"Task failed: {e}")
+    # Task is automatically released back to pending queue
+```
+
+## Stale Task Recovery
+
+TaskCore automatically reclaims stale tasks (default: 4 hours timeout):
+
+```python
+# Custom timeout in seconds
+queue.queue.reclaim_stale_tasks(timeout_seconds=3600)  # 1 hour
+```
+
+## License
+
+MIT License - see [LICENSE](LICENSE) file for details.

taskcore_lib-0.0.2.dist-info/RECORD

@@ -0,0 +1,9 @@
+taskcore/__init__.py,sha256=5qiCW0vnEOZmiTQEEahEGLtOWJOv3Ors6xoEXqwV6io,43
+taskcore/base.py,sha256=hZlc2Tq7nU2UcxnoJs20-2iH-zQXBT3KdNFetWG1lr4,8019
+taskcore/cli.py,sha256=uC7nhjxglDRi-SaxpFfZCgzP1_WkHwPvhyYPWZqxlGI,1533
+taskcore/dist.py,sha256=6qaFc85YsbgFbvCgBCII7OYeP368f0m4IKF-pF-LLog,1692
+taskcore_lib-0.0.2.dist-info/licenses/LICENSE,sha256=ldUsph09cc1l_w5lD-IeZQR_RmAr_IfCq1m1sK7timk,1070
+taskcore_lib-0.0.2.dist-info/METADATA,sha256=3CuXnfjFcLRbqPgGpHfkIToLrtjpoO0DBQ-FG1bazUc,6164
+taskcore_lib-0.0.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+taskcore_lib-0.0.2.dist-info/top_level.txt,sha256=SYLIoRvq8EDnjU5FdNfJldfSFwgHcfwij8xqu3bREAg,9
+taskcore_lib-0.0.2.dist-info/RECORD,,

taskcore_lib-0.0.2.dist-info/WHEEL

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

taskcore_lib-0.0.2.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Gonçalo Faria
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

taskcore_lib-0.0.2.dist-info/top_level.txt

@@ -0,0 +1 @@
+taskcore