Inject unified Tunix and vLLM environment fingerprinting metadata

PiperOrigin-RevId: 915141032

Author: sizhit2

tests/sft/metrics_logger_test.py

@@ -226,8 +226,27 @@ def test_backend_kwargs_are_passed_to_backends(self):
           resume="must",
           id="12345",
           settings=mock_wandb.Settings.return_value,
+          config={
+              "tunix_version": mock.ANY,
+              "tunix_commit": mock.ANY,
+              "vllm_version": mock.ANY,
+              "vllm_commit": mock.ANY,
+          },
       )
 
+  def test_get_module_info(self):
+    """Tests environment fingerprinting module info extraction."""
+    ver, commit = metrics_logger._get_module_info(None)
+    self.assertEqual(ver, "not_installed")
+    self.assertEqual(commit, "not_installed")
+
+    dummy_module = mock.Mock()
+    dummy_module.__version__ = "0.4.0+dev.deadbeef"
+    del dummy_module.__file__
+    ver, commit = metrics_logger._get_module_info(dummy_module)
+    self.assertEqual(ver, "0.4.0+dev.deadbeef")
+    self.assertEqual(commit, "deadbeef")
+
 
 if __name__ == "__main__":
   absltest.main()

tunix/sft/metrics_logger.py

@@ -3,6 +3,9 @@
 import collections
 import dataclasses
 import enum
+import importlib.metadata
+import os
+import subprocess
 from typing import Any, Callable
 
 from absl import logging
@@ -11,6 +14,41 @@
 import numpy as np
 from tunix.utils import env_utils
 
+try:
+  import vllm  # pylint: disable=g-import-not-at-top
+except ImportError:
+  vllm = None
+
+
+def _get_module_info(
+    module: Any, path_for_git: str = "", package_name: str = ""
+) -> tuple[str, str]:
+  """Resolves (version, commit) concisely for environment fingerprinting."""
+  if module is None and not package_name and not path_for_git:
+    return "not_installed", "not_installed"
+
+  ver = getattr(module, "__version__", "")
+  if package_name:
+    try:
+      ver = importlib.metadata.version(package_name)
+    except Exception:
+      pass
+  ver = ver or "unknown"
+
+  commit = ""
+  path = path_for_git or getattr(module, "__file__", "")
+  if path and os.path.exists(path):
+    try:
+      d = os.path.dirname(path) if os.path.isfile(path) else path
+      commit = subprocess.check_output(
+          ["git", "rev-parse", "HEAD"], cwd=d, stderr=subprocess.DEVNULL, text=True
+      ).strip()
+    except Exception:
+      pass
+
+  commit = commit or (ver.split("+")[-1] if "+" in ver else "unknown")
+  return ver, commit
+
 LoggingBackend = metrax_logging.LoggingBackend
 TensorboardBackend = metrax_logging.TensorboardBackend
 WandbBackend = metrax_logging.WandbBackend
@@ -55,6 +93,16 @@ def create_backends(self) -> list[LoggingBackend]:
     if jax.process_index() != 0:
       return []
 
+    tunix_version, tunix_commit = _get_module_info(
+        None, path_for_git=__file__, package_name="google-tunix"
+    )
+    vllm_version, vllm_commit = _get_module_info(vllm)
+
+    logging.info("=== Tunix Environment Fingerprint ===")
+    logging.info("Tunix: version=%s, commit=%s", tunix_version, tunix_commit)
+    logging.info("vLLM:  version=%s, commit=%s", vllm_version, vllm_commit)
+    logging.info("=====================================")
+
     # Case 1: Override. Use user-provided factories.
     if (
         "custom_backend" in self.backend_kwargs
@@ -83,7 +131,13 @@ def create_backends(self) -> list[LoggingBackend]:
           )
       )
       try:
-        wandb_kwargs = kwargs_dict.get("wandb", {})
+        wandb_kwargs = dict(kwargs_dict.get("wandb", {}))
+        wandb_config = dict(wandb_kwargs.get("config", {}))
+        wandb_config["tunix_version"] = tunix_version
+        wandb_config["tunix_commit"] = tunix_commit
+        wandb_config["vllm_version"] = vllm_version
+        wandb_config["vllm_commit"] = vllm_commit
+        wandb_kwargs["config"] = wandb_config
         active_backends.append(
             WandbBackend(
                 project=self.project_name,