🔧 chore: optimize dependency configuration and improve logging

Author: zrr1999

.gitignore

@@ -6,3 +6,6 @@ dist/
 .cursorrules
 
 __version__.py
+
+.coverage
+coverage.xml

nanoflow/config.py

@@ -9,6 +9,22 @@ class TaskConfig(BaseModel):
 
 
 class WorkflowConfig(BaseModel):
+    """
+    Workflow config file.
+
+    Example:
+    >>> config = WorkflowConfig(
+    >>>     name="test",
+    >>>     tasks={
+    >>>         "task1": TaskConfig(command="echo 'task1'"),
+    >>>         "task2": TaskConfig(command="echo 'task2'", deps=["task1"]),
+    >>>         "task3": TaskConfig(command="echo 'task3'", deps=["task2"]),
+    >>>     }
+    >>> )
+    >>> config.to_nodes()
+    >>> {'task1': [], 'task2': ['task1'], 'task3': ['task2']}
+    """
+
     name: str
     tasks: dict[str, TaskConfig]
 

nanoflow/task.py

@@ -4,6 +4,7 @@
 from collections.abc import Callable
 from typing import Any, Generic, ParamSpec, TypeVar, overload
 
+from loguru import logger
 from pydantic import BaseModel, ConfigDict
 
 from .resource_pool import ResourcePool
@@ -26,6 +27,7 @@ def submit(self, *args: P.args, **kwargs: P.kwargs) -> asyncio.Task[R]:
         async def wrapper_fn():
             if self.resource_pool is not None:
                 resource = await self.resource_pool.acquire()
+                logger.info(f"Acquired resource: {resource}")
                 if self.resource_modifier is not None:
                     fn = self.resource_modifier(self.fn, resource)
                 else:
@@ -34,6 +36,7 @@ async def wrapper_fn():
                     return await asyncio.to_thread(fn, *args, **kwargs)
                 finally:
                     self.resource_pool.release(resource)
+                    logger.info(f"Released resource: {resource}")
             else:
                 return await asyncio.to_thread(self.fn, *args, **kwargs)
 

nanoflow/utils.py

@@ -54,14 +54,14 @@ def get_available_gpus(threshold=0.05) -> list[int]:
 @workflow
 async def execute_parallel_tasks(config: WorkflowConfig):
     def create_gpu_task(command: str) -> Task:
-        environ = os.environ.copy()
-
         def set_visible_gpu(fn: Callable[[], bytes], resource: int) -> Callable[[], bytes]:
+            environ = os.environ.copy()
             environ["CUDA_VISIBLE_DEVICES"] = str(resource)
 
             def inner_fn() -> bytes:
                 res = subprocess.run(command, shell=True, env=environ)
                 if res.returncode != 0:
+                    # TODO: raise an exception
                     logger.error(f"Command failed with return code {res.returncode}")
                     return res.stderr
                 return res.stdout

pyproject.toml

@@ -5,17 +5,20 @@ description = "A simple and efficient workflow framework for Python"
 authors = [{ name = "Zhan Rongrui", email = "[email protected]" }]
 dependencies = [
     "loguru>=0.7.0",
-    "matplotlib>=3.9.0",
     "networkx>=3.0",
     "pydantic>=2.0.0",
-    "rich>=13.8.0",
     "toml>=0.10.0",
+    "rich>=13.8.0",
     "typer>=0.12.5",
 ]
 requires-python = ">=3.11"
 readme = "README.md"
 license = { text = "MIT" }
 
+[project.optional-dependencies]
+server = ["fastapi>=0.115.0", "uvicorn>=0.30.6"]
+plot = ["matplotlib>=3.9.0"]
+
 [project.urls]
 Homepage = "https://github.com/zrr1999/nanoflow.git"
 Repository = "https://github.com/zrr1999/nanoflow.git"
@@ -78,6 +81,14 @@ write_template = "__version__ = '{}'"
 includes = ["nanoflow"]
 source-includes = ["README.md"]
 
+[tool.uv]
+dev-dependencies = [
+    "xdoctest",
+    "pytest",
+    "pytest-cov",
+    "coverage",
+]
+
 [build-system]
 requires = ["pdm-backend"]
 build-backend = "pdm.backend"