Files
language-conditioned-psycho…/research/environment/smoke_test.py
T

146 lines
4.3 KiB
Python

import importlib
import json
import os
from pathlib import Path
import sys
import traceback
PROJECT_ROOT = Path(__file__).resolve().parents[2]
CACHE_ROOT = PROJECT_ROOT / ".cache"
cache_settings = {
"LOCALAPPDATA": CACHE_ROOT / "localappdata",
"MPLCONFIGDIR": CACHE_ROOT / "matplotlib",
"NUMBA_CACHE_DIR": CACHE_ROOT / "numba",
"SKB_DATA_DIRECTORY": CACHE_ROOT / "skrub",
"HF_HOME": CACHE_ROOT / "huggingface",
"TORCH_HOME": CACHE_ROOT / "torch",
"XDG_CACHE_HOME": CACHE_ROOT,
}
for key, path in cache_settings.items():
path.mkdir(parents=True, exist_ok=True)
os.environ[key] = str(path)
os.environ["PYTENSOR_FLAGS"] = f"base_compiledir={CACHE_ROOT / 'pytensor'}"
import numpy as np
results = {"python": sys.version, "imports": {}, "tests": {}}
def record(name, fn):
try:
value = fn()
results["tests"][name] = {"status": "passed", "detail": value}
except Exception as exc:
results["tests"][name] = {
"status": "failed",
"detail": f"{type(exc).__name__}: {exc}",
"traceback": traceback.format_exc(),
}
for module_name in [
"numpy", "scipy", "pandas", "polars", "pyarrow", "duckdb",
"statsmodels", "pymc", "pytensor", "arviz", "nutpie",
"jax", "numpyro", "blackjax", "bambi", "semopy",
"factor_analyzer", "pingouin", "girth", "torch", "torchvision",
"transformers", "sentence_transformers", "xgboost",
"tabpfn",
]:
try:
module = importlib.import_module(module_name)
results["imports"][module_name] = {
"status": "passed",
"version": getattr(module, "__version__", "unknown"),
}
except Exception as exc:
results["imports"][module_name] = {
"status": "failed",
"detail": f"{type(exc).__name__}: {exc}",
}
def pymc_test():
import pymc as pm
rng = np.random.default_rng(20260920)
observed = rng.normal(0.35, 1.0, size=40)
with pm.Model():
mu = pm.Normal("mu", 0.0, 1.0)
pm.Normal("y", mu, 1.0, observed=observed)
inference = pm.sample(
draws=50,
tune=50,
chains=1,
cores=1,
random_seed=20260920,
progressbar=False,
compute_convergence_checks=False,
)
posterior_mean = float(inference.posterior["mu"].mean())
return {"posterior_mean": posterior_mean, "draws": 50, "chains": 1}
def sem_test():
import pandas as pd
import semopy
rng = np.random.default_rng(20260920)
latent = rng.normal(size=160)
frame = pd.DataFrame(
{
"x1": 0.8 * latent + rng.normal(scale=0.4, size=160),
"x2": 0.7 * latent + rng.normal(scale=0.5, size=160),
"x3": 0.9 * latent + rng.normal(scale=0.3, size=160),
}
)
model = semopy.Model("factor =~ x1 + x2 + x3")
fit = model.fit(frame)
return {"success": bool(fit.success), "objective": float(fit.fun)}
def factor_test():
from factor_analyzer import FactorAnalyzer
rng = np.random.default_rng(20260920)
latent = rng.normal(size=(180, 1))
data = latent @ np.array([[0.8, 0.7, 0.9, 0.6]]) + rng.normal(
scale=0.45, size=(180, 4)
)
fitted = FactorAnalyzer(n_factors=1, rotation=None).fit(data)
return {"loading_shape": list(fitted.loadings_.shape)}
def torch_test():
import torch
if not torch.cuda.is_available():
raise RuntimeError("torch.cuda.is_available() is false")
left = torch.randn((256, 256), device="cuda")
right = torch.randn((256, 256), device="cuda")
result = left @ right
torch.cuda.synchronize()
return {
"torch": torch.__version__,
"cuda_runtime": torch.version.cuda,
"device": torch.cuda.get_device_name(0),
"finite": bool(torch.isfinite(result).all().item()),
}
record("pymc_sampling", pymc_test)
record("semopy_cfa", sem_test)
record("factor_analyzer", factor_test)
record("torch_cuda", torch_test)
serialized = json.dumps(results, ensure_ascii=False, indent=2)
(Path(__file__).parent / "smoke-test-result.json").write_text(
serialized, encoding="utf-8"
)
print(serialized)
failed_imports = [k for k, v in results["imports"].items() if v["status"] != "passed"]
failed_tests = [k for k, v in results["tests"].items() if v["status"] != "passed"]
raise SystemExit(1 if failed_imports or failed_tests else 0)