fix(bazel): adjust test setup and environment handling for Bazel compatibility

* Enable to see projects
* bazel compliance: Rewrite collect_env_info and direct cache a tempfile

Author: georgeyiasemis

MODULE.bazel

@@ -23,28 +23,28 @@ bazel_dep(name = "aspect_rules_py", version = "1.3.2")
 # This allows sharding of tests and easier setup
 archive_override(
     module_name = "aspect_rules_py",
-    integrity = "sha256-p7Fo6yMoh96SurOcGt1uLt95vRf/RUeuPvr/1oA6jzw=",
-    strip_prefix = "rules_py-1296c6b88156594543b7f41ceb331bdc9a053055",
-    urls = ["https://github.com/aspect-build/rules_py/archive/1296c6b88156594543b7f41ceb331bdc9a053055.zip"],
+    integrity = "sha256-0EJJ2KgKvlISrxVk5Tmc74iMMJNZJZ/3dG55azcc8sA=",
+    strip_prefix = "rules_py-1445ccaf3665cb5d8f78da4f5fc4d73fd36fa165",
+    urls = ["https://github.com/aspect-build/rules_py/archive/1445ccaf3665cb5d8f78da4f5fc4d73fd36fa165.zip"],
 )
 
 bazel_dep(name = "rules_uv", version = "0.65.0", dev_dependency = True)
 
 python = use_extension("@rules_python//python/extensions:python.bzl", "python")
 python.toolchain(
     is_default = True,
-    python_version = "3.12",
+    python_version = "3.11",
 )
 
 pip = use_extension("@rules_python//python/extensions:pip.bzl", "pip")
 pip.parse(
     enable_implicit_namespace_pkgs = True,
-    hub_name = "pip",
-    python_version = "3.12",
+    hub_name = "direct_pip",
+    python_version = "3.11",
     requirements_darwin = ":requirements_darwin.txt",
     requirements_lock = ":requirements_linux.txt",
 )
-use_repo(pip, "pip")
+use_repo(pip, "direct_pip")
 
 npm = use_extension("@aspect_rules_js//npm:extensions.bzl", "npm")
 npm.npm_translate_lock(

direct/BUILD.bazel

@@ -1,39 +1,39 @@
 load("@aspect_rules_py//py:defs.bzl", "py_binary", "py_library")
-load("@pip//:requirements.bzl", "requirement")
+load("@direct_pip//:requirements.bzl", "requirement")
 load("@rules_cc//cc:defs.bzl", "cc_library")
 load("//tools:cython_rules.bzl", "pyx_library")
 
 cc_library(
     name = "numpy_headers",
     hdrs = [
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/__multiarray_api.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/__ufunc_api.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/_dtype_api.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/_neighborhood_iterator_imp.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/_numpyconfig.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/arrayobject.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/arrayscalars.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/experimental_dtype_api.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/halffloat.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/ndarrayobject.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/ndarraytypes.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/noprefix.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/npy_1_7_deprecated_api.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/npy_3kcompat.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/npy_common.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/npy_cpu.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/npy_endian.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/npy_interrupt.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/npy_math.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/npy_no_deprecated_api.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/npy_os.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/numpyconfig.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/old_defines.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/random/bitgen.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/random/distributions.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/random/libdivide.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/ufuncobject.h",
-        "@@rules_python~~pip~pip_312_numpy//:site-packages/numpy/core/include/numpy/utils.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/__multiarray_api.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/__ufunc_api.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/_dtype_api.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/_neighborhood_iterator_imp.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/_numpyconfig.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/arrayobject.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/arrayscalars.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/experimental_dtype_api.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/halffloat.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/ndarrayobject.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/ndarraytypes.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/noprefix.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/npy_1_7_deprecated_api.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/npy_3kcompat.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/npy_common.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/npy_cpu.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/npy_endian.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/npy_interrupt.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/npy_math.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/npy_no_deprecated_api.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/npy_os.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/numpyconfig.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/old_defines.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/random/bitgen.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/random/distributions.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/random/libdivide.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/ufuncobject.h",
+        "@@rules_python~~pip~direct_pip_311_numpy//:site-packages/numpy/core/include/numpy/utils.h",
     ],
     strip_include_prefix = "/site-packages/numpy/core/include",
     visibility = ["//visibility:private"],
@@ -63,7 +63,10 @@ pyx_library(
 py_library(
     name = "algorithms",
     srcs = glob(["algorithms/**/*.py"]),
-    deps = [requirement("torch"), requirement("numpy")],
+    deps = [
+        requirement("torch"),
+        requirement("numpy"),
+    ],
 )
 
 py_library(
@@ -130,7 +133,9 @@ py_library(
 
 py_library(
     name = "direct_lib",
-    srcs = glob(["*.py"]),
+    srcs = glob(["**/*.py"]),
+    imports = ["."],
+    visibility = ["//visibility:public"],
     deps = [
         ":algorithms",
         ":common",
@@ -145,14 +150,16 @@ py_library(
         requirement("torchvision"),
         requirement("omegaconf"),
     ],
-    visibility = ["//visibility:public"],
 )
 
 py_binary(
     name = "direct",
     srcs = glob([
         "cli/**/*.py",
     ]),
+    data = [
+        "//projects:all_configs",
+    ],
     main = "cli/cli.py",
     deps = [
         ":direct_lib",

direct/environment.py

@@ -12,10 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import argparse
+import importlib.metadata
 import logging
 import os
 import pathlib
+import platform
 import sys
+import tempfile
 from collections import namedtuple
 from typing import Callable, Dict, Optional, Tuple, Union
 
@@ -33,10 +36,120 @@
 
 # Environmental variables
 DIRECT_ROOT_DIR = pathlib.Path(pathlib.Path(__file__).resolve().parent.parent)
-DIRECT_CACHE_DIR = pathlib.Path(os.environ.get("DIRECT_CACHE_DIR", str(DIRECT_ROOT_DIR)))
-DIRECT_MODEL_DOWNLOAD_DIR = (
-    pathlib.Path(os.environ.get("DIRECT_MODEL_DOWNLOAD_DIR", str(DIRECT_ROOT_DIR))) / "downloaded_models"
-)
+
+
+def resolve_cache_dir() -> pathlib.Path:
+    cache_dir_path = pathlib.Path(os.environ.get("DIRECT_CACHE_DIR", str(DIRECT_ROOT_DIR)))
+    # Check if the directory is writable
+    if os.access(str(cache_dir_path), os.W_OK):
+        logger.info("Using cache directory: %s", cache_dir_path)
+        return cache_dir_path
+    if "DIRECT_CACHE_DIR" in os.environ:
+        env_path = pathlib.Path(os.environ["DIRECT_CACHE_DIR"])
+        if os.access(str(env_path), os.W_OK):
+            logger.info("Using cache directory: %s", env_path)
+            return env_path
+    try:
+        tmpdir = os.environ.get("TMPDIR", tempfile.gettempdir())
+        cache_dir = pathlib.Path(tmpdir) / "direct_cache"
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        if os.access(str(cache_dir), os.W_OK):
+            logger.info("Using cache directory: %s", cache_dir)
+            return cache_dir
+    except Exception as e:
+        logger.warning("Failed to create or access cache directory in TMPDIR: %s", e)
+
+    # Fallback to a default tmp directory
+    fallback = pathlib.Path(tempfile.gettempdir()) / "direct_cache"
+    fallback.mkdir(parents=True, exist_ok=True)
+    logger.warning("Falling back to cache directory: %s", fallback)
+    return fallback
+
+
+DIRECT_CACHE_DIR = resolve_cache_dir()
+DIRECT_MODEL_DOWNLOAD_DIR = DIRECT_CACHE_DIR / "downloaded_models"
+
+
+def collect_env_info() -> str:
+    """Collects environment information.
+
+    Returns
+    -------
+    env_info: str
+        Environment information as a formatted string.
+    """
+    SystemEnv = namedtuple(
+        "SystemEnv",
+        [
+            "torch_version",
+            "is_debug_build",
+            "cuda_compiled_version",
+            "python_version",
+            "python_platform",
+            "os",
+            "libc_version",
+            "is_cuda_available",
+            "cuda_runtime_version",
+            "cudnn_version",
+            "pip_packages",
+            "cpu_info",
+        ],
+    )
+
+    def safe_version(pkg):
+        try:
+            return importlib.metadata.version(pkg)
+        except importlib.metadata.PackageNotFoundError:
+            return "Not installed"
+
+    def get_cudnn_version():
+        try:
+            return str(torch.backends.cudnn.version()) if torch.backends.cudnn.is_available() else "Unavailable"
+        except Exception:
+            return "Unknown"
+
+    def get_cpu_info():
+        try:
+            return platform.processor() or platform.machine()
+        except Exception:
+            return "Unknown"
+
+    pip_packages = {pkg: safe_version(pkg) for pkg in ["torch", "numpy", "triton", "optree", "mypy", "flake8", "onnx"]}
+    pip_str = "\n    " + "\n    ".join(f"{pkg}=={ver}" for pkg, ver in pip_packages.items())
+
+    def pretty_print(env):
+        lines = [
+            f"PyTorch version: {env.torch_version}",
+            f"Is debug build: {env.is_debug_build}",
+            f"CUDA used to build PyTorch: {env.cuda_compiled_version}",
+            f"Python version: {env.python_version}",
+            f"Python platform: {env.python_platform}",
+            f"OS: {env.os}",
+            f"Libc version: {env.libc_version}",
+            f"Is CUDA available: {env.is_cuda_available}",
+            f"CUDA runtime version: {env.cuda_runtime_version}",
+            f"cuDNN version: {env.cudnn_version}",
+            f"CPU info: {env.cpu_info}",
+            f"Relevant pip packages: {env.pip_packages}",
+        ]
+        return "\n" + "\n".join(lines)
+
+    return pretty_print(
+        SystemEnv(
+            torch_version=torch.__version__,
+            is_debug_build=str(getattr(torch.version, "debug", "Unknown")),
+            cuda_compiled_version=getattr(torch.version, "cuda", "None"),
+            python_version=sys.version.replace("\n", " "),
+            python_platform=platform.platform(),
+            os=platform.platform(),
+            libc_version="-".join(platform.libc_ver()) if sys.platform.startswith("linux") else "N/A",
+            is_cuda_available=str(torch.cuda.is_available()),
+            cuda_runtime_version=getattr(torch.version, "cuda", "No CUDA"),
+            cudnn_version=get_cudnn_version(),
+            pip_packages=pip_str,
+            cpu_info=get_cpu_info(),
+        )
+    )
 
 
 def load_model_config_from_name(model_name: str) -> Callable:
@@ -152,7 +265,7 @@ def setup_logging(
     logger.info("Run name: %s", run_name)
     logger.info("Config file: %s", cfg_filename)
     logger.info("CUDA %s - cuDNN %s", torch.version.cuda, torch.backends.cudnn.version())
-    logger.info("Environment information: %s", collect_env.get_pretty_env_info())
+    logger.info("Environment information: %s", collect_env_info())
     logger.info("DIRECT version: %s", direct.__version__)
     git_hash = direct.utils.git_hash()
     logger.info("Git hash: %s", git_hash if git_hash else "N/A")

projects/BUILD.bazel

@@ -0,0 +1,5 @@
+filegroup(
+    name = "all_configs",
+    srcs = glob(["**/*.yaml"]),
+    visibility = ["//visibility:public"],
+)

requirements_darwin.txt

@@ -778,9 +778,7 @@ scipy==1.15.2 \
 setuptools==79.0.0 \
     --hash=sha256:9828422e7541213b0aacb6e10bbf9dd8febeaa45a48570e09b6d100e063fc9f9 \
     --hash=sha256:b9ab3a104bedb292323f53797b00864e10e434a3ab3906813a7169e4745b912a
-    # via
-    #   tensorboard
-    #   torch
+    # via tensorboard
 six==1.17.0 \
     --hash=sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274 \
     --hash=sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81

requirements_linux.txt

@@ -848,9 +848,7 @@ scipy==1.15.2 \
 setuptools==79.0.0 \
     --hash=sha256:9828422e7541213b0aacb6e10bbf9dd8febeaa45a48570e09b6d100e063fc9f9 \
     --hash=sha256:b9ab3a104bedb292323f53797b00864e10e434a3ab3906813a7169e4745b912a
-    # via
-    #   tensorboard
-    #   torch
+    # via tensorboard
 six==1.17.0 \
     --hash=sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274 \
     --hash=sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81

tests/BUILD.bazel

@@ -1,5 +1,5 @@
 load("@aspect_rules_py//py:defs.bzl", "py_test")
-load("@pip//:requirements.bzl", "requirement")
+load("@direct_pip//:requirements.bzl", "requirement")
 
 REQUIREMENTS = [
     "//direct:direct_lib",

tests/test_cli/BUILD.bazel

@@ -1,5 +1,5 @@
 load("@aspect_rules_py//py:defs.bzl", "py_test")
-load("@pip//:requirements.bzl", "requirement")
+load("@direct_pip//:requirements.bzl", "requirement")
 
 REQUIREMENTS = [
     "//direct:direct_lib",

tests/tests_common/BUILD.bazel

@@ -1,5 +1,5 @@
 load("@aspect_rules_py//py:defs.bzl", "py_test")
-load("@pip//:requirements.bzl", "requirement")
+load("@direct_pip//:requirements.bzl", "requirement")
 
 REQUIREMENTS = [
     "//direct:direct_lib",

tests/tests_data/BUILD.bazel

@@ -1,5 +1,5 @@
 load("@aspect_rules_py//py:defs.bzl", "py_test")
-load("@pip//:requirements.bzl", "requirement")
+load("@direct_pip//:requirements.bzl", "requirement")
 
 REQUIREMENTS = [
     "//direct:direct_lib",