fix pyprojec.toml and pre-commit wandb issues (#712)

.pre-commit-config.yaml

@@ -38,4 +38,4 @@ repos:
     hooks:
     -   id: mypy
         args: [--ignore-missing-imports]
-        additional_dependencies: [wandb, types-PyYAML]
+        additional_dependencies: [wandb==0.17.8, types-PyYAML]

pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
     "tokenizers>=0.15.2",
     "transformers>=4.41.2",
     "optax>=0.1.9",
-    "wandb>=0.16.6,<0.17.6",
+    "wandb>=0.17.8",
     "scipy<=1.12.0",
     "draccus>=0.8.0",
     "pyarrow>=11.0.0",
@@ -105,13 +105,6 @@ test = [
     "pytest-asyncio",
 ]
 
-#[tool.setuptools.packages.find]
-#where = ["src"]
-#include = ["levanter", "levanter.*"]
-
-
-[tool.setuptools]
-packages = ["levanter"]
-
-[tool.setuptools.package-dir]
-levanter = "src/levanter"
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["levanter", "levanter.*"]

src/levanter/tracker/wandb.py

@@ -182,9 +182,9 @@ def init(self, run_id: Optional[str]) -> WandbTracker:
             if wandb.run is not None:
                 wandb.run.log_artifact(str(requirements_path), name="requirements.txt", type="requirements")
 
-        wandb.summary["num_devices"] = jax.device_count()
-        wandb.summary["num_hosts"] = jax.process_count()
-        wandb.summary["backend"] = jax.default_backend()
+        wandb.summary["num_devices"] = jax.device_count()  # type: ignore
+        wandb.summary["num_hosts"] = jax.process_count()  # type: ignore
+        wandb.summary["backend"] = jax.default_backend()  # type: ignore
 
         return WandbTracker(r)