wandb-util 0.1.0__tar.gz

wandb_util-0.1.0/PKG-INFO

@@ -0,0 +1,11 @@
+Metadata-Version: 2.3
+Name: wandb-util
+Version: 0.1.0
+Summary: Utility tool providing extended functionality for W&B (Weights & Biases)
+Author: Xiao Li
+Author-email: Xiao Li <xiaoli3397@gmail.com>
+Requires-Dist: click>=8.3.3
+Requires-Dist: wandb>=0.26.1
+Requires-Python: >=3.13
+Description-Content-Type: text/markdown
+

wandb_util-0.1.0/pyproject.toml

@@ -0,0 +1,20 @@
+[project]
+name = "wandb-util"
+version = "0.1.0"
+description = "Utility tool providing extended functionality for W&B (Weights & Biases)"
+readme = "README.md"
+authors = [
+    { name = "Xiao Li", email = "xiaoli3397@gmail.com" }
+]
+requires-python = ">=3.13"
+dependencies = [
+    "click>=8.3.3",
+    "wandb>=0.26.1",
+]
+
+[project.scripts]
+wandb-util = "wandb_util.__main__:cli"
+
+[build-system]
+requires = ["uv_build>=0.11.7,<0.12.0"]
+build-backend = "uv_build"

wandb_util-0.1.0/src/wandb_util/__init__.py

@@ -0,0 +1,2 @@
+def hello() -> str:
+    return "Hello from wandb-util!"

wandb_util-0.1.0/src/wandb_util/__main__.py

@@ -0,0 +1,23 @@
+import click
+
+from .commands import artifact, log, run
+
+
+@click.group()
+@click.option("--entity", "-e", required=True, help="W&B entity (username or team).")
+@click.option("--project", "-p", required=True, help="W&B project name.")
+@click.pass_context
+def cli(ctx, entity, project):
+    """WandB utility tool providing extended functionality."""
+    ctx.ensure_object(dict)
+    ctx.obj["entity"] = entity
+    ctx.obj["project"] = project
+
+
+cli.add_command(artifact)
+cli.add_command(log)
+cli.add_command(run)
+
+
+if __name__ == "__main__":
+    cli()

wandb_util-0.1.0/src/wandb_util/commands/__init__.py

@@ -0,0 +1,5 @@
+from .artifact import artifact
+from .log import log
+from .run import run
+
+__all__ = ["artifact", "log", "run"]

wandb_util-0.1.0/src/wandb_util/commands/artifact.py

@@ -0,0 +1,88 @@
+from pathlib import Path
+
+import click
+import wandb
+
+from ..utils import handle_api_errors
+
+
+def fmt_size(n_bytes: int) -> str:
+    size: float = float(n_bytes)
+    for unit in ("B", "KB", "MB", "GB"):
+        if size < 1024:
+            return f"{size:.1f} {unit}"
+        size /= 1024
+    return f"{size:.1f} TB"
+
+
+def _list_artifacts(
+    api: wandb.Api, entity: str, project: str, run_id: str, artifact_type: str
+) -> None:
+    """Core logic to list artifacts for a run."""
+    run = api.run(f"{entity}/{project}/{run_id}")
+    artifacts = sorted(
+        [a for a in run.logged_artifacts() if a.type == artifact_type],
+        key=lambda x: x.metadata.get("global_step", 0),
+    )
+
+    if not artifacts:
+        print(f"No artifacts found for run '{run_id}'")
+        return
+
+    print(f"Run: {run.name} ({run_id})  state={run.state}")
+    print()
+    print(f"{'Version':<28} {'Step':>8} {'Size':>10}  {'Created'}")
+    print("-" * 72)
+    for a in artifacts:
+        step = a.metadata.get("global_step", "-")
+        created = a.created_at[:19].replace("T", " ") if a.created_at else "-"
+        print(f"{a.name:<28} {str(step):>8} {fmt_size(a.size):>10}  {created}")
+
+
+def _download_artifact(
+    api: wandb.Api, entity: str, project: str, artifact_path: str, output_dir: str
+) -> None:
+    """Core logic to download an artifact."""
+    artifact = api.artifact(f"{entity}/{project}/{artifact_path}", type="model")
+    out = Path(output_dir)
+    click.echo(f"Downloading {artifact.name} ({fmt_size(artifact.size)}) -> {out}/")
+    artifact.download(root=str(out))
+    click.echo("Done.")
+
+
+@click.group()
+@click.pass_context
+def artifact(ctx):
+    """Manage W&B artifacts."""
+    pass
+
+
+@artifact.command()
+@click.argument("run_id")
+@click.option("--type", "-t", default="model", help="Artifact type (default: model).")
+@click.pass_context
+@handle_api_errors
+def list(ctx, run_id, type):
+    """List artifacts for a run."""
+    entity = ctx.obj["entity"]
+    project = ctx.obj["project"]
+    api = wandb.Api()
+    _list_artifacts(api, entity, project, run_id, type)
+
+
+artifact.add_command(list, name="ls")
+
+
+@artifact.command()
+@click.argument("artifact_path")
+@click.option(
+    "--output-dir", "-o", default=".", help="Directory to download into (default: .)."
+)
+@click.pass_context
+@handle_api_errors
+def download(ctx, artifact_path, output_dir):
+    """Download an artifact."""
+    entity = ctx.obj["entity"]
+    project = ctx.obj["project"]
+    api = wandb.Api()
+    _download_artifact(api, entity, project, artifact_path, output_dir)

wandb_util-0.1.0/src/wandb_util/commands/log.py

@@ -0,0 +1,88 @@
+import click
+import wandb
+
+from ..utils import handle_api_errors
+
+
+def _list_logs(
+    api: wandb.Api,
+    entity: str,
+    project: str,
+    run_id: str,
+    show_history: bool = False,
+    filter_keys: str | None = None,
+    prefix: str | None = None,
+    last: int = 10,
+) -> None:
+    """Core logic to list logs for a run."""
+    run = api.run(f"{entity}/{project}/{run_id}")
+
+    print(f"Run:     {run.name} ({run.id})")
+    print(f"State:   {run.state}")
+    print(f"Project: {entity}/{project}")
+    print()
+
+    filter_set = set(filter_keys.split(",")) if filter_keys else None
+
+    if show_history:
+        history = run.history()
+        if filter_set:
+            cols = [c for c in history.columns if c in filter_set or c == "_step"]
+        else:
+            cols = list(history.columns)
+
+        df = history[cols].tail(last)
+        print(df.to_string(index=False))
+    else:
+        summary = run.summary.items()
+        if prefix:
+            summary = {k: v for k, v in summary if k.startswith(prefix)}
+        else:
+            summary = dict(summary)
+        if filter_set:
+            summary = {k: v for k, v in summary.items() if k in filter_set}
+
+        if not summary:
+            print("No metrics found in summary.")
+            return
+
+        key_width = max(len(k) for k in summary)
+        print(f"{'Metric':<{key_width}}  Value")
+        print("-" * (key_width + 20))
+        for key, value in sorted(summary.items()):
+            if isinstance(value, float):
+                print(f"{key:<{key_width}}  {value:.6g}")
+            else:
+                print(f"{key:<{key_width}}  {value}")
+
+
+@click.group()
+@click.pass_context
+def log(ctx):
+    """Manage W&B logs."""
+    pass
+
+
+@log.command()
+@click.argument("run_id")
+@click.option("--history", is_flag=True, help="Show history instead of summary.")
+@click.option("--prefix", help="Filter summary by key prefix (e.g., eval/, train/).")
+@click.option("--keys", "-k", help="Filter by comma-separated keys.")
+@click.option(
+    "--last",
+    "-l",
+    default=10,
+    type=int,
+    help="Show last N history entries (default: 10).",
+)
+@click.pass_context
+@handle_api_errors
+def list(ctx, run_id, history, prefix, keys, last):
+    """List logs for a run."""
+    entity = ctx.obj["entity"]
+    project = ctx.obj["project"]
+    api = wandb.Api()
+    _list_logs(api, entity, project, run_id, history, keys, prefix, last)
+
+
+log.add_command(list, name="ls")

wandb_util-0.1.0/src/wandb_util/commands/run.py

@@ -0,0 +1,47 @@
+import click
+import wandb
+
+from ..utils import handle_api_errors
+
+
+def _list_runs(api: wandb.Api, entity: str, project: str, limit: int) -> None:
+    """Core logic to list runs."""
+    runs = api.runs(f"{entity}/{project}", order="-created_at")
+    print(f"{'ID':<12} {'State':<10} {'Name':<40} {'Created':<20} {'Steps':>8}")
+    print("-" * 96)
+    count = 0
+    for run in runs:
+        if count >= limit:
+            break
+        created = run.created_at[:19].replace("T", " ") if run.created_at else "-"
+        steps = run.summary.get("_step", "-")
+        print(
+            f"{run.id:<12} {run.state:<10} {run.name[:40]:<40} {created:<20} {str(steps):>8}"
+        )
+        count += 1
+    if count == 0:
+        print(f"No runs found in '{entity}/{project}'")
+
+
+@click.group()
+@click.pass_context
+def run(ctx):
+    """Manage W&B runs."""
+    pass
+
+
+@run.command()
+@click.option(
+    "--limit", "-l", default=10, type=int, help="Maximum number of runs to display."
+)
+@click.pass_context
+@handle_api_errors
+def list(ctx, limit):
+    """List runs."""
+    entity = ctx.obj["entity"]
+    project = ctx.obj["project"]
+    api = wandb.Api()
+    _list_runs(api, entity, project, limit)
+
+
+run.add_command(list, name="ls")

wandb_util-0.1.0/src/wandb_util/utils.py

@@ -0,0 +1,30 @@
+import functools
+
+import click
+import wandb
+
+
+def handle_api_errors(func):
+    """Decorator to handle common wandb API errors."""
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except wandb.errors.CommError as e:
+            click.secho(
+                f"Error: Failed to connect to wandb API. {str(e)}", fg="red", err=True
+            )
+            raise SystemExit(1)
+        except wandb.errors.AuthenticationError:
+            click.secho(
+                "Error: Authentication failed. Check your API key in ~/.netrc",
+                fg="red",
+                err=True,
+            )
+            raise SystemExit(1)
+        except Exception as e:
+            click.secho(f"Error: {str(e)}", fg="red", err=True)
+            raise SystemExit(1)
+
+    return wrapper