Lazyfying

src/curate_questions/create_question_set/main.py

@@ -15,14 +15,14 @@
 import json
 import logging
 import os
-import random
 import sys
 from collections.abc import Callable
 from copy import deepcopy
 from datetime import datetime, timedelta
 from enum import Enum
 from fractions import Fraction
 
+import numpy as np
 import pandas as pd
 from tqdm import tqdm
 from utils import gcp
@@ -94,15 +94,19 @@ def process_questions(
     single_generation_func: Callable,
     show_plots: bool,
     question_set_target: QuestionSetTarget,
+    random_state=None,
 ) -> dict:
     """Sample from `questions` to get the number of questions needed.
 
     Args:
         questions (dict): Source questions keyed by source name, each with a "dfq" DataFrame
         to_questions (dict): Allocation info keyed by source with "num_questions_to_sample"
-        single_generation_func (Callable): Sampling function taking (values, n) and returning DataFrame
+        single_generation_func (Callable): Sampling function taking (values, n, random_state) and
+            returning a DataFrame
         show_plots (bool): Whether to display distribution plots
         question_set_target (QuestionSetTarget): Target question set ("llm" or "human")
+        random_state: Seed/``np.random.RandomState`` threaded to ``single_generation_func`` for
+            reproducibility. ``None`` (default) samples without a fixed seed.
 
     Returns
         processed_questions (dict): Deep copy of questions with sampled DataFrames
@@ -118,7 +122,7 @@ def process_questions(
         df_available = values["dfq"].copy()
 
         # Sample questions for this source
-        values["dfq"] = single_generation_func(values, num_single)
+        values["dfq"] = single_generation_func(values, num_single, random_state)
         df_sampled = values["dfq"]
         num_found += len(df_sampled)
 
@@ -179,20 +183,21 @@ def process_questions(
     return processed_questions
 
 
-def human_sample_questions(values: dict, n_single: int) -> pd.DataFrame:
+def human_sample_questions(values: dict, n_single: int, random_state=None) -> pd.DataFrame:
     """Get questions for the human question set by sampling from LLM questions.
 
     Args:
         values (dict): Source data dict containing "dfq" DataFrame
         n_single (int): Number of questions to sample
+        random_state: Seed/``np.random.RandomState`` for reproducible sampling (anything
+            ``DataFrame.sample`` accepts). ``None`` (default) samples without a fixed seed.
 
     Returns
         dfq (pd.DataFrame): Randomly sampled questions
     """
     dfq = values["dfq"].copy()
-    indices_to_sample_from = dfq.index.tolist()
-    indices = random.sample(indices_to_sample_from, min(n_single, len(indices_to_sample_from)))
-    return dfq.loc[indices]
+    n = min(n_single, len(dfq))
+    return dfq.sample(n=n, random_state=random_state)
 
 
 def get_bin_label(bin_config: dict, bin_type: str) -> str:
@@ -743,18 +748,36 @@ def add_line_chart(
     fig.show()
 
 
-def stratified_sample_questions(dfq: pd.DataFrame, n_target: int) -> pd.DataFrame:
+def _as_random_state(random_state):
+    """Normalize a seed/RandomState/None to a RandomState (or None).
+
+    Returns ``None`` unchanged (preserving unseeded behavior) and a ``RandomState`` unchanged, but
+    promotes a bare int seed to a ``RandomState`` so successive ``.sample`` calls in a loop
+    *decorrelate* (advancing one shared generator) instead of reusing the same seed per iteration.
+    """
+    if random_state is None or isinstance(random_state, np.random.RandomState):
+        return random_state
+    return np.random.RandomState(random_state)
+
+
+def stratified_sample_questions(
+    dfq: pd.DataFrame, n_target: int, random_state=None
+) -> pd.DataFrame:
     """Sample questions using stratified sampling to achieve target distribution.
 
     This ensures we get the desired distribution regardless of source data skew.
 
     Args:
         dfq (pd.DataFrame): DataFrame with bin_weight column and composite bins
         n_target (int): Number of questions to sample
+        random_state: Seed/``np.random.RandomState`` for reproducible per-bin sampling (anything
+            ``DataFrame.sample`` accepts). ``None`` (default) samples without a fixed seed. Thread a
+            single ``RandomState`` instance through to decorrelate the successive per-bin draws.
 
     Returns
         result (pd.DataFrame): Sampled questions
     """
+    random_state = _as_random_state(random_state)
     if len(dfq) == 0 or n_target == 0:
         return pd.DataFrame()
 
@@ -811,22 +834,24 @@ def stratified_sample_questions(dfq: pd.DataFrame, n_target: int) -> pd.DataFram
     for bin_name, n_samples in bin_samples.items():
         if n_samples > 0:
             bin_df = dfq_weighted[dfq_weighted["composite_bin"] == bin_name]
-            sampled = bin_df.sample(n=n_samples, replace=False)
+            sampled = bin_df.sample(n=n_samples, replace=False, random_state=random_state)
             sampled_dfs.append(sampled)
 
     if not sampled_dfs:
         raise ValueError("Stratified sampling produced no results.")
     return pd.concat(sampled_dfs, ignore_index=True)
 
 
-def sample_market_questions(dfq: pd.DataFrame, n_target: int) -> pd.DataFrame:
+def sample_market_questions(dfq: pd.DataFrame, n_target: int, random_state=None) -> pd.DataFrame:
     """Sample market questions using multi-dimensional binning strategy.
 
     Ensures balanced sampling across market probability values and time horizons.
 
     Args:
         dfq (pd.DataFrame): Market questions
         n_target (int): Number of questions to sample
+        random_state: Seed/``np.random.RandomState`` threaded to ``stratified_sample_questions``
+            for reproducibility. ``None`` (default) samples without a fixed seed.
 
     Returns
         df_result (pd.DataFrame): Sampled questions
@@ -840,6 +865,7 @@ def sample_market_questions(dfq: pd.DataFrame, n_target: int) -> pd.DataFrame:
     df_result = stratified_sample_questions(
         dfq=dfq,
         n_target=n_target,
+        random_state=random_state,
     )
     df_result = df_result.drop(
         columns=[
@@ -852,7 +878,7 @@ def sample_market_questions(dfq: pd.DataFrame, n_target: int) -> pd.DataFrame:
     return df_result
 
 
-def llm_sample_questions(values: dict, n_single: int) -> pd.DataFrame:
+def llm_sample_questions(values: dict, n_single: int, random_state=None) -> pd.DataFrame:
     """Generate questions for the LLM question set.
 
     For market questions: Sample using binning strategy.
@@ -861,23 +887,26 @@ def llm_sample_questions(values: dict, n_single: int) -> pd.DataFrame:
     Args:
         values (dict): Source data dict containing "dfq" DataFrame
         n_single (int): Number of questions to sample
+        random_state: Seed/``np.random.RandomState`` threaded to the market/category samplers for
+            reproducibility. ``None`` (default) samples without a fixed seed.
 
     Returns
         df (pd.DataFrame): Sampled questions
     """
     dfq = values["dfq"].copy()
     source = dfq["source"].iloc[0]
+    random_state = _as_random_state(random_state)
 
     if source in question_curation.MARKET_SOURCES:
         # Use binning-based sampling for market questions
-        return sample_market_questions(dfq, n_single)
+        return sample_market_questions(dfq, n_single, random_state=random_state)
     else:
         # Use existing category-based sampling for data sources
         allocation = allocate_across_categories(num_questions=n_single, dfq=dfq)
 
         dfs = []
         for key, value in allocation.items():
-            dfs.append(dfq[dfq["category"] == key].sample(value))
+            dfs.append(dfq[dfq["category"] == key].sample(value, random_state=random_state))
         return pd.concat(dfs, ignore_index=True)
 
 
@@ -1206,6 +1235,12 @@ def driver(_: None) -> None:
         )
         HUMAN_QUESTIONS.update(human_questions_of_question_type)
 
+    # Optional reproducibility: when QUESTION_SET_SEED is set, thread a single RandomState through
+    # both sampling passes so the published set is deterministic. Unset (the default) preserves the
+    # historical unseeded behaviour.
+    seed_env = os.getenv("QUESTION_SET_SEED")
+    random_state = np.random.RandomState(int(seed_env)) if seed_env not in (None, "") else None
+
     # Sample questions
     logger.info("LLM SET")
     LLM_QUESTIONS = process_questions(
@@ -1214,6 +1249,7 @@ def driver(_: None) -> None:
         single_generation_func=llm_sample_questions,
         show_plots=env.RUNNING_LOCALLY,
         question_set_target=QuestionSetTarget.LLM,
+        random_state=random_state,
     )
 
     logger.info("HUMAN SET")
@@ -1223,6 +1259,7 @@ def driver(_: None) -> None:
         single_generation_func=human_sample_questions,
         show_plots=False,
         question_set_target=QuestionSetTarget.HUMAN,
+        random_state=random_state,
     )
 
     write_questions(LLM_QUESTIONS, question_set_target=QuestionSetTarget.LLM)

src/helpers/env.py

@@ -1,25 +1,47 @@
-"""Environment variables."""
+"""Environment variables.
+
+Values are read from ``os.environ`` lazily on each attribute access (PEP 562 module
+``__getattr__``) rather than snapshotted at import.
+"""
 
 import os
 
-PROJECT_ID = os.environ.get("CLOUD_PROJECT")
-CLOUD_DEPLOY_REGION = os.environ.get("CLOUD_DEPLOY_REGION")
-QUESTION_BANK_BUCKET = os.environ.get("QUESTION_BANK_BUCKET")
-QUESTION_SETS_BUCKET = os.environ.get("QUESTION_SETS_BUCKET")
-FORECAST_SETS_BUCKET = os.environ.get("FORECAST_SETS_BUCKET")
-FORECAST_SETS_TRANSCRIPTS_BUCKET = os.environ.get("FORECAST_SETS_TRANSCRIPTS_BUCKET")
-PROCESSED_FORECAST_SETS_BUCKET = os.environ.get("PROCESSED_FORECAST_SETS_BUCKET")
-PUBLIC_RELEASE_BUCKET = os.environ.get("PUBLIC_RELEASE_BUCKET")
-WEBSITE_BUCKET = os.environ.get("WEBSITE_BUCKET")
-WEBSITE_STAGING_ASSETS_BUCKET = os.environ.get("WEBSITE_STAGING_ASSETS_BUCKET")
-LLM_BASELINE_DOCKER_IMAGE_NAME = os.environ.get("LLM_BASELINE_DOCKER_IMAGE_NAME")
-LLM_BASELINE_DOCKER_REPO_NAME = os.environ.get("LLM_BASELINE_DOCKER_REPO_NAME")
-LLM_BASELINE_PUB_SUB_TOPIC_NAME = os.environ.get("LLM_BASELINE_PUB_SUB_TOPIC_NAME")
-LLM_BASELINE_STAGING_BUCKET = os.environ.get("LLM_BASELINE_STAGING_BUCKET")
-LLM_BASELINE_SERVICE_ACCOUNT = os.environ.get("LLM_BASELINE_SERVICE_ACCOUNT")
-LLM_BASELINE_STAGING_BUCKET = os.environ.get("LLM_BASELINE_STAGING_BUCKET")
-LLM_BASELINE_NEWS_BUCKET = os.environ.get("LLM_BASELINE_NEWS_BUCKET")
-NUM_CPUS = int(os.environ.get("NUM_CPUS", 1))
-RUNNING_LOCALLY = bool(int(os.environ.get("RUNNING_LOCALLY", False)))
-BUCKET_MOUNT_POINT = os.environ.get("BUCKET_MOUNT_POINT", "")
-WORKSPACE_BUCKET = os.environ.get("WORKSPACE_BUCKET")
+# Names read as plain strings (``None`` if unset).
+_STR_VARS = {
+    "PROJECT_ID": "CLOUD_PROJECT",
+    "CLOUD_DEPLOY_REGION": "CLOUD_DEPLOY_REGION",
+    "QUESTION_BANK_BUCKET": "QUESTION_BANK_BUCKET",
+    "QUESTION_SETS_BUCKET": "QUESTION_SETS_BUCKET",
+    "FORECAST_SETS_BUCKET": "FORECAST_SETS_BUCKET",
+    "FORECAST_SETS_TRANSCRIPTS_BUCKET": "FORECAST_SETS_TRANSCRIPTS_BUCKET",
+    "PROCESSED_FORECAST_SETS_BUCKET": "PROCESSED_FORECAST_SETS_BUCKET",
+    "PUBLIC_RELEASE_BUCKET": "PUBLIC_RELEASE_BUCKET",
+    "WEBSITE_BUCKET": "WEBSITE_BUCKET",
+    "WEBSITE_STAGING_ASSETS_BUCKET": "WEBSITE_STAGING_ASSETS_BUCKET",
+    "LLM_BASELINE_DOCKER_IMAGE_NAME": "LLM_BASELINE_DOCKER_IMAGE_NAME",
+    "LLM_BASELINE_DOCKER_REPO_NAME": "LLM_BASELINE_DOCKER_REPO_NAME",
+    "LLM_BASELINE_PUB_SUB_TOPIC_NAME": "LLM_BASELINE_PUB_SUB_TOPIC_NAME",
+    "LLM_BASELINE_STAGING_BUCKET": "LLM_BASELINE_STAGING_BUCKET",
+    "LLM_BASELINE_SERVICE_ACCOUNT": "LLM_BASELINE_SERVICE_ACCOUNT",
+    "LLM_BASELINE_NEWS_BUCKET": "LLM_BASELINE_NEWS_BUCKET",
+    "WORKSPACE_BUCKET": "WORKSPACE_BUCKET",
+}
+
+
+def __getattr__(name):
+    """Read environment variables lazily on each access (PEP 562)."""
+    if name in _STR_VARS:
+        return os.environ.get(_STR_VARS[name])
+    if name == "NUM_CPUS":
+        return int(os.environ.get("NUM_CPUS", 1))
+    if name == "RUNNING_LOCALLY":
+        return bool(int(os.environ.get("RUNNING_LOCALLY", False)))
+    if name == "BUCKET_MOUNT_POINT":
+        return os.environ.get("BUCKET_MOUNT_POINT", "")
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__():
+    """Expose the lazily-read environment variable names to ``dir()``/autocomplete."""
+    extra = {"NUM_CPUS", "RUNNING_LOCALLY", "BUCKET_MOUNT_POINT"}
+    return sorted(set(globals()) | set(_STR_VARS) | extra)

src/helpers/keys.py

@@ -1,9 +1,39 @@
-"""utils for key-related tasks in llm-benchmark."""
+"""utils for key-related tasks in llm-benchmark.
+
+Secrets are resolved lazily on first attribute access (PEP 562 module ``__getattr__``)
+and memoized. This ensures that merely *importing* this module performs no network/Secret Manager call.
+"""
 
 from google.cloud import secretmanager
 
 from . import env
 
+# Public attribute name -> Secret Manager secret name. The two differ only for GOOGLE/GEMINI.
+_SECRET_NAMES = {
+    # LLM
+    "API_KEY_ANTHROPIC": "API_KEY_ANTHROPIC",
+    "API_KEY_OPENAI": "API_KEY_OPENAI",
+    "API_KEY_TOGETHERAI": "API_KEY_TOGETHERAI",
+    "API_KEY_GOOGLE": "API_KEY_GEMINI",
+    "API_KEY_MISTRAL": "API_KEY_MISTRAL",
+    "API_KEY_XAI": "API_KEY_XAI",
+    # QUESTION DATASET SOURCES
+    "API_EMAIL_ACLED": "API_EMAIL_ACLED",
+    "API_PASSWORD_ACLED": "API_PASSWORD_ACLED",
+    "API_KEY_FRED": "API_KEY_FRED",
+    # QUESTION MARKET SOURCES
+    "API_KEY_METACULUS": "API_KEY_METACULUS",
+    "API_KEY_POLYMARKET": "API_KEY_POLYMARKET",
+    "API_KEY_INFER": "API_KEY_INFER",
+    # WORKFLOW BOT
+    "API_SLACK_BOT_NOTIFICATION": "API_SLACK_BOT_NOTIFICATION",
+    "API_SLACK_BOT_CHANNEL": "API_SLACK_BOT_CHANNEL",
+    # GITHUB
+    "API_GITHUB_DATASET_REPO_URL": "API_GITHUB_DATASET_REPO_URL",
+}
+
+_cache: dict = {}
+
 
 def get_secret(secret_name, version_id="latest"):
     """
@@ -27,19 +57,16 @@ def get_secret_that_may_not_exist(secret_name, version_id="latest"):
         return None
 
 
-# QUESTION DATASET SOURCES
-API_EMAIL_ACLED = get_secret(secret_name="API_EMAIL_ACLED")
-API_PASSWORD_ACLED = get_secret(secret_name="API_PASSWORD_ACLED")
-API_KEY_FRED = get_secret("API_KEY_FRED")
-
-# QUESTION MARKET SOURCES
-API_KEY_METACULUS = get_secret(secret_name="API_KEY_METACULUS")
-API_KEY_POLYMARKET = get_secret("API_KEY_POLYMARKET")
-API_KEY_INFER = get_secret("API_KEY_INFER")
+def __getattr__(name):
+    """Lazily resolve and memoize ``API_*`` secrets on first access (PEP 562)."""
+    secret_name = _SECRET_NAMES.get(name)
+    if secret_name is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    if name not in _cache:
+        _cache[name] = get_secret(secret_name)
+    return _cache[name]
 
-# WORKFLOW BOT
-API_SLACK_BOT_NOTIFICATION = get_secret(secret_name="API_SLACK_BOT_NOTIFICATION")
-API_SLACK_BOT_CHANNEL = get_secret(secret_name="API_SLACK_BOT_CHANNEL")
 
-# GITHUB
-API_GITHUB_DATASET_REPO_URL = get_secret(secret_name="API_GITHUB_DATASET_REPO_URL")
+def __dir__():
+    """Expose the lazily-resolved secret names to ``dir()``/autocomplete."""
+    return sorted(set(globals()) | set(_SECRET_NAMES))