Spaces:
Sleeping
Sleeping
File size: 4,356 Bytes
3e12d11 ec6e9a7 3e12d11 ec6e9a7 3e12d11 ec6e9a7 3e12d11 ec6e9a7 3e12d11 ec6e9a7 3e12d11 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """Data manager for loading and caching SCB employment pipeline results.
This module encapsulates the logic for computing the SCB-only
transformations in ``pipeline.py`` and persisting the result to disk.
It adds a small amount of resilience around caching and uses
``logging`` instead of printing directly to stdout. The cache file
includes a version tag to make it easy to invalidate caches when
fundamental changes are made to the pipeline logic.
"""
import os
import tempfile
import logging
from pathlib import Path
from functools import lru_cache
import pandas as pd
from . import pipeline
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Cache setup
# ---------------------------------------------------------------------------
# A version tag to embed into the cache filenames. Bump this value
# whenever the underlying ``pipeline`` logic changes in a way that
# invalidates existing caches (or when you want to force a fresh recompute).
CACHE_VERSION: str = "v1"
def _resolve_cache_dir() -> Path:
"""Select a writable directory for caching.
The lookup order is:
1. The ``DATA_CACHE_DIR`` environment variable, if set.
2. A ``data`` folder at the repository root.
3. A temporary directory in ``/tmp``.
Each candidate path is tested for writability by attempting to
create and delete a sentinel file. The first path that succeeds
is returned. If none succeed, a final fallback directory in ``/tmp``
is created and returned.
"""
candidates: list[Path] = []
env = os.getenv("DATA_CACHE_DIR")
if env:
# Expand relative or user paths to absolute
candidates.append(Path(env).expanduser().resolve())
# Repo root /data (two levels up from this file)
candidates.append(Path(__file__).resolve().parent.parent / "data")
# Temp fallback
candidates.append(Path(tempfile.gettempdir()) / "employment_ai_cache")
for path in candidates:
try:
path.mkdir(parents=True, exist_ok=True)
test_file = path / ".write_test"
test_file.write_text("ok", encoding="utf-8")
test_file.unlink()
return path
except Exception:
continue
# Final fallback: ensure the last candidate exists
fallback = Path(tempfile.gettempdir()) / "employment_ai_cache"
fallback.mkdir(parents=True, exist_ok=True)
return fallback
# Resolve the directory once at import time
DATA_DIR: Path = _resolve_cache_dir()
# Single cache file for the SCB-only output DataFrame.
SCB_CACHE: Path = DATA_DIR / f"scb_employment_{CACHE_VERSION}.csv"
def _atomic_to_csv(df: pd.DataFrame, path: Path) -> None:
"""Write a DataFrame to CSV atomically.
The CSV is first written to a temporary file in the same directory
and then renamed to the final location. This avoids leaving a
partially written file if the process is interrupted mid‑write.
"""
path.parent.mkdir(parents=True, exist_ok=True)
tmp_path = path.with_suffix(path.suffix + ".tmp")
df.to_csv(tmp_path, index=False)
tmp_path.replace(path)
@lru_cache(maxsize=1)
def _compute_pipeline_payload() -> pd.DataFrame:
"""Run the SCB-only pipeline calculation (cached in-process)."""
return pipeline.run_pipeline()
def load_payload() -> pd.DataFrame:
"""Load employment data from disk cache if available, otherwise compute and save.
To force a rebuild, delete the cache file or bump `CACHE_VERSION` and restart.
"""
if SCB_CACHE.exists():
logger.info("Loading pipeline output from cache directory %s", DATA_DIR)
try:
return pd.read_csv(SCB_CACHE)
except Exception as exc:
# If reading the cache fails, fall back to recomputing
logger.warning(
"Error reading cache file %s: %s; falling back to recompute",
SCB_CACHE,
exc,
)
logger.info("Computing SCB employment data – this may take a while…")
payload = _compute_pipeline_payload()
# Persist to disk atomically
try:
_atomic_to_csv(payload, SCB_CACHE)
logger.info("Cache updated: %s", SCB_CACHE.name)
except Exception as exc:
logger.warning("Could not write cache file: %s", exc)
return payload
|