mouse-experiment 0.1.0__py3-none-any.whl

mouse/runner/__main__.py

@@ -0,0 +1,20 @@
+import os
+import sys
+from pathlib import Path
+
+
+def main() -> None:
+    from mouse.runner.experiment import main as run_main
+
+    config_path = ""
+    if len(sys.argv) > 1:
+        arg = sys.argv[1]
+        if not os.path.isabs(arg) and "/" not in arg and "\\" not in arg:
+            arg = str(Path.cwd() / "configs" / arg)
+        config_path = arg
+
+    run_main(config_path)
+
+
+if __name__ == "__main__":
+    main()

mouse/runner/auth.py

@@ -0,0 +1,111 @@
+"""Authentication helpers for third-party services (Hugging Face, Weights & Biases, Trackio)."""
+
+import os
+from typing import Any
+from huggingface_hub import login as hf_login
+
+_hf_authenticated: bool = False
+_wandb_authenticated: bool = False
+_trackio_initialized: bool = False
+
+
+def mirror_hf_token_env() -> None:
+    """Copy alternate HF token env vars into ``HF_TOKEN`` for hub clients.
+
+    ``transformers`` / ``huggingface_hub`` read ``HF_TOKEN`` from the environment
+    during ``AutoConfig.from_pretrained`` and similar calls. Those run from
+    ``load_config()`` *before* ``setup_huggingface()`` in ``loop.main``, so
+    ``RUNPOD_HF_TOKEN`` must be mirrored up front.
+    """
+    if os.environ.get("HF_TOKEN"):
+        return
+    for env_var in ("RUNPOD_HF_TOKEN",):
+        token = os.environ.get(env_var)
+        if token:
+            os.environ["HF_TOKEN"] = token
+            break
+
+
+def setup_huggingface() -> None:
+    """Authenticate with Hugging Face using env vars.
+
+    Expects one of: HF_TOKEN, RUNPOD_HF_TOKEN.
+    Idempotent: the /whoami-v2 verification is performed at most once per process
+    to avoid hitting the strict HF rate limit.
+    """
+    global _hf_authenticated
+    if _hf_authenticated:
+        return
+
+    hf_token_envs = ["HF_TOKEN", "RUNPOD_HF_TOKEN"]
+    hf_token = None
+    for env_var in hf_token_envs:
+        hf_token = os.getenv(env_var)
+        if hf_token:
+            break
+    if not hf_token:
+        raise ValueError(f"None of the following environment variables are set: {hf_token_envs}")
+
+    os.environ["HF_TOKEN"] = hf_token
+    hf_login(token=hf_token)
+    _hf_authenticated = True
+
+
+def setup_wandb() -> None:
+    """Authenticate with Weights & Biases using env vars.
+
+    Expects one of: WANDB_TOKEN, RUNPOD_WANDB_TOKEN.
+    Idempotent: login is performed at most once per process.
+    """
+    global _wandb_authenticated
+    if _wandb_authenticated:
+        return
+
+    import wandb  # type: ignore[import-untyped]
+
+    wandb_token_envs = ["WANDB_TOKEN", "RUNPOD_WANDB_TOKEN"]
+    wandb_token = None
+    for env_var in wandb_token_envs:
+        wandb_token = os.getenv(env_var)
+        if wandb_token:
+            break
+    if not wandb_token:
+        raise ValueError(f"None of the following environment variables are set: {wandb_token_envs}")
+
+    wandb.login(key=wandb_token, verify=True)
+    _wandb_authenticated = True
+
+
+def setup_trackio(project: str, run_name: str, config: dict, space_id: str | None = None) -> Any:
+    """Initialize a Trackio run for experiment tracking.
+
+    Trackio stores metrics locally in SQLite and optionally syncs to a Hugging
+    Face Space dashboard. It is a drop-in complement to W&B — both can log the
+    same metrics in parallel.
+
+    Args:
+        project: Trackio project name (groups related runs).
+        run_name: Unique name for this run.
+        config: Flat or nested dict of hyperparameters to attach to the run.
+        space_id: Optional HF Space to sync the dashboard to, e.g. ``"user/my-space"``.
+            Also read from the ``TRACKIO_SPACE_ID`` environment variable when not given.
+
+    Returns:
+        The active trackio run object (supports ``.log()`` and ``.finish()``).
+    """
+    global _trackio_initialized
+    import trackio  # type: ignore[import-untyped]
+
+    init_kwargs: dict = dict(
+        project=project,
+        name=run_name,
+        config=config,
+    )
+    if space_id:
+        init_kwargs["space_id"] = space_id
+    elif os.environ.get("TRACKIO_SPACE_ID"):
+        init_kwargs["space_id"] = os.environ["TRACKIO_SPACE_ID"]
+
+    run = trackio.init(**init_kwargs)
+    _trackio_initialized = True
+    return run