Merge pull request #475 from switchbox-data/474-factor-out-aws_storage_option-out-of-eia-scripts

alexhyunminlee · web-flow · commit 5a7960836154 · 2026-06-30T10:52:03.000-04:00
Refactor get aws storage options function
diff --git a/.gitignore b/.gitignore
@@ -41,6 +41,7 @@ data/isone/capacity/ara/parquet/
 data/isone/capacity/fca/parquet/
 data/isone/lmp/parquet/
 data/isone/hourly_demand/zones/
+data/isone/hourly_demand/utilities/
 data/isone/ancillary/parquet/
 data/nyiso/icap/parquet/
 data/nyiso/lbmp/zips/
diff --git a/AGENTS.md b/AGENTS.md
@@ -81,7 +81,7 @@ After individual CAIRO runs complete, post-processing scripts consolidate result
 | **`rate_design/`**                  | Package root. Heat pump rate design lives under `rate_design/hp_rates/`.                                                                                                                                                                                                                                                                                                                                                                                            |
 | **`rate_design/hp_rates/`**         | Shared scenario entrypoint (`run_scenario.py`), shared **Justfile** (primary task interface for all states). State-specific thin Justfiles and `config/` dirs live under `rate_design/hp_rates/{ny,ri}/`.                                                                                                                                                                                                                                                           |
 | **`rate_design/hp_rates/{ny,ri}/`** | State-specific thin Justfile (imports shared), `state.env`, and `config/` (tariffs JSON in tariffs/electric and tariffs/gas, tariff_maps CSV in tariff_maps/electric and tariff_maps/gas, marginal_costs). Large artifacts (buildstock raw/processed, cairo_cases) are git-ignored; sync via S3 or keep local.                                                                                                                                                      |
-| **`data/eia/hourly_loads/`**        | EIA zone load fetch and utility load aggregation; eia_region_config (state/utility config, get_aws_storage_options); Justfile for fetch-zone-data and aggregate-utility-loads.                                                                                                                                                                                                                                                                                      |
+| **`data/eia/hourly_loads/`**        | EIA zone load fetch and utility load aggregation; eia_region_config (state/utility config); Justfile for fetch-zone-data and aggregate-utility-loads.                                                                                                                                                                                                                                                                                                               |
 | **`data/eia/861/`**                 | EIA-861 utility stats (PUDL yearly sales); fetch_electric_utility_stat_parquets.py; Justfile build-utility-stats (local parquet), update (upload to s3://data.sb/eia/861/electric_utility_stats/), fetch-utility-stats STATE (CSV to stdout).                                                                                                                                                                                                                       |
 | **`data/fred/cpi/`**                | FRED CPI series; Justfile fetch-cpi (local parquet/), upload (sync to s3://data.sb/fred/cpi/).                                                                                                                                                                                                                                                                                                                                                                      |
 | **`data/aspe/fpl/`**                | ASPE Federal Poverty Guidelines fetch; Justfile fetch. Output: utils/post/data/fpl_guidelines.yaml (used by LMI discount logic).                                                                                                                                                                                                                                                                                                                                    |
diff --git a/context/code/cairo/cairo_elastic_cluster.md b/context/code/cairo/cairo_elastic_cluster.md
@@ -40,7 +40,7 @@ With a **Dask Distributed** cluster, the driver and scheduler run on one node; w
   - **Overhead:** You must deploy workers (e.g. EC2 or Batch) that have the same mount. With Fargate or generic Batch workers, that can be fiddly (custom AMI or entrypoint that mounts S3).
 
 - **B. S3 URIs + storage_options (no shared mount)**\
-  Pass **S3 URIs** (e.g. `s3://data.sb/nrel/resstock/...`) and **storage_options** (e.g. `{"region_name": "us-west-2"}` or from `get_aws_storage_options()`) through the stack so that workers can call `pd.read_parquet(s3_uri, storage_options=...)` and read directly from S3.
+  Pass **S3 URIs** (e.g. `s3://data.sb/nrel/resstock/...`) and **storage_options** (e.g. `{"region_name": "us-west-2"}` or from `utils.file_io.get_aws_storage_options()`) through the stack so that workers can call `pd.read_parquet(s3_uri, storage_options=...)` and read directly from S3.
   - **Overhead:** Requires code changes in both codebases (see below). No need for a shared mount; any worker with AWS credentials and network can read.
 
 Option **B** is more portable (works with Fargate, Batch, any node with S3 access) and is the one that forces the concrete changes below. Option A is “no code change” but constrains how you run workers.
@@ -92,7 +92,7 @@ CAIRO does **not** need to know about “the cluster”; it only needs to be abl
 2. **Pass storage_options into CAIRO**
    - Where the platform calls CAIRO (e.g. `_return_load`, `return_buildingstock`, and the top-level `simulate()` or helpers that take paths), it must pass `storage_options` when the paths are S3. That implies:
      - CAIRO’s public API (e.g. `MeetRevenueSufficiencySystemWide`, `simulate()`, or the load/buildingstock functions) must accept an optional `storage_options` (or a “run config” that includes it) and pass it through to all parquet-reading code.
-     - The platform already has `get_aws_storage_options()` (e.g. in `data.eia.hourly_loads.eia_region_config` or similar). When running in “cluster mode” with S3 URIs, the platform should pass that (or equivalent) into CAIRO.
+     - The platform already has `get_aws_storage_options()` (in `utils.file_io`). When running in “cluster mode” with S3 URIs, the platform should pass that (or equivalent) into CAIRO.
 
 3. **Use Dask Distributed instead of the process scheduler**
    - Instead of `dask.config.set(scheduler="processes", num_workers=...)`, the platform (or a cluster launcher script) would create a `distributed.Client` connected to a Dask scheduler (e.g. one started by dask-cloudprovider), and either:
diff --git a/data/eia/hourly_loads/eia_region_config.py b/data/eia/hourly_loads/eia_region_config.py
@@ -5,7 +5,6 @@
 - utility service areas used by utility aggregation scripts
 """
 
-import os
 from dataclasses import dataclass
 from typing import TypedDict
 
@@ -139,14 +138,3 @@ def get_utility_zone_mapping_for_state(state: str) -> dict[str, list[str]]:
                 break
 
     return utility_zone_mapping
-
-
-def get_aws_storage_options() -> dict[str, str]:
-    """Return Polars-compatible AWS storage options for S3 access."""
-    aws_region = (
-        os.getenv("AWS_DEFAULT_REGION") or os.getenv("AWS_REGION") or "us-west-2"
-    )
-    return {
-        "region": aws_region,
-        "default_region": aws_region,
-    }
diff --git a/utils/cairo.py b/utils/cairo.py
@@ -14,7 +14,7 @@
 from cairo.rates_tool.loads import __timeshift__
 from cloudpathlib import S3Path
 
-from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+from utils.file_io import get_aws_storage_options
 from utils.types import ElectricUtility
 
 CambiumPathLike = str | Path | S3Path
diff --git a/utils/data_prep/marginal_costs/generate_bulk_tx_mc.py b/utils/data_prep/marginal_costs/generate_bulk_tx_mc.py
@@ -40,7 +40,7 @@
 
 from dotenv import load_dotenv
 
-from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+from utils.file_io import get_aws_storage_options
 from utils.data_prep.marginal_costs.bulk_tx_isone import (
     AESC_2024_AVOIDED_PTF_KW_YEAR,
     DEFAULT_N_PEAK_HOURS,
diff --git a/utils/data_prep/marginal_costs/generate_supply_ancillary_mc.py b/utils/data_prep/marginal_costs/generate_supply_ancillary_mc.py
@@ -120,7 +120,7 @@ def main() -> None:
     )
 
     load_dotenv()
-    from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+    from utils.file_io import get_aws_storage_options
 
     storage_options = get_aws_storage_options()
 
diff --git a/utils/data_prep/marginal_costs/generate_supply_capacity_mc.py b/utils/data_prep/marginal_costs/generate_supply_capacity_mc.py
@@ -6,7 +6,7 @@
 
 from dotenv import load_dotenv
 
-from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+from utils.file_io import get_aws_storage_options
 from utils.data_prep.marginal_costs.supply_capacity_nyiso import (
     N_PEAK_HOURS_PER_MONTH,
     compute_supply_capacity_mc,
diff --git a/utils/data_prep/marginal_costs/generate_supply_energy_mc.py b/utils/data_prep/marginal_costs/generate_supply_energy_mc.py
@@ -6,7 +6,7 @@
 
 from dotenv import load_dotenv
 
-from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+from utils.file_io import get_aws_storage_options
 from utils.data_prep.marginal_costs.supply_energy import (
     compute_isone_supply_energy_mc,
     compute_pjm_supply_energy_mc,
diff --git a/utils/data_prep/marginal_costs/generate_utility_tx_dx_mc.py b/utils/data_prep/marginal_costs/generate_utility_tx_dx_mc.py
@@ -60,10 +60,8 @@
 from cloudpathlib import S3Path
 from dotenv import load_dotenv
 
-from data.eia.hourly_loads.eia_region_config import (
-    get_aws_storage_options,
-    get_state_config,
-)
+from data.eia.hourly_loads.eia_region_config import get_state_config
+from utils.file_io import get_aws_storage_options
 from utils.data_prep.marginal_costs.supply_utils import (
     warn_if_multiple_partition_parquets,
 )
diff --git a/utils/file_io.py b/utils/file_io.py
@@ -4,13 +4,25 @@
 
 import io
 import itertools
+import os
 from pathlib import Path
 from typing import Any
 
 import boto3
 import polars as pl
 
 
+def get_aws_storage_options() -> dict[str, str]:
+    """Return Polars-compatible AWS storage options for S3 access."""
+    aws_region = (
+        os.getenv("AWS_DEFAULT_REGION") or os.getenv("AWS_REGION") or "us-west-2"
+    )
+    return {
+        "region": aws_region,
+        "default_region": aws_region,
+    }
+
+
 def read_csv_from_s3(s3_path: str, **kwargs: Any) -> pl.DataFrame:
     """Read a CSV file from S3 into a Polars DataFrame.
 
diff --git a/utils/mid/compute_fair_default_inputs.py b/utils/mid/compute_fair_default_inputs.py
@@ -15,7 +15,7 @@
 from cloudpathlib import S3Path
 from dotenv import load_dotenv
 
-from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+from utils.file_io import get_aws_storage_options
 from utils.loads import ELECTRIC_LOAD_COL, ELECTRIC_PV_COL, grid_consumption_expr
 from utils.loads import scan_resstock_loads, scan_resstock_loads_monthly
 from utils.mid.compute_subclass_rr import (
diff --git a/utils/mid/compute_flat_discount_inputs.py b/utils/mid/compute_flat_discount_inputs.py
@@ -9,7 +9,7 @@
 from cloudpathlib import S3Path
 from dotenv import load_dotenv
 
-from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+from utils.file_io import get_aws_storage_options
 from utils.mid.compute_subclass_rr import (
     BAT_METRIC_CHOICES,
     DEFAULT_BAT_METRIC,
diff --git a/utils/mid/compute_seasonal_discount_inputs.py b/utils/mid/compute_seasonal_discount_inputs.py
@@ -9,7 +9,7 @@
 from cloudpathlib import S3Path
 from dotenv import load_dotenv
 
-from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+from utils.file_io import get_aws_storage_options
 from utils.mid.compute_subclass_rr import (
     DEFAULT_BAT_METRIC,
     DEFAULT_GROUP_COL,
diff --git a/utils/mid/compute_subclass_rr.py b/utils/mid/compute_subclass_rr.py
@@ -23,7 +23,7 @@
 from cloudpathlib import S3Path
 from dotenv import load_dotenv
 
-from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+from utils.file_io import get_aws_storage_options
 from utils.loads import ELECTRIC_PV_COL, grid_consumption_expr, scan_resstock_loads
 from utils.pre.season_config import (
     DEFAULT_SEASONAL_DISCOUNT_WINTER_MONTHS,
diff --git a/utils/mid/prepare_fair_default_tariffs.py b/utils/mid/prepare_fair_default_tariffs.py
@@ -23,7 +23,7 @@
 from cloudpathlib import S3Path
 from dotenv import load_dotenv
 
-from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+from utils.file_io import get_aws_storage_options
 from utils import get_aws_region
 from utils.cairo import (
     add_bulk_tx_and_dist_and_sub_tx_marginal_cost,
diff --git a/utils/post/apply_ny_lmi_discounts_to_cairo_bills.py b/utils/post/apply_ny_lmi_discounts_to_cairo_bills.py
@@ -29,7 +29,7 @@
 from cloudpathlib import S3Path
 from dotenv import load_dotenv
 
-from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+from utils.file_io import get_aws_storage_options
 from utils.post.lmi_common import (
     assign_ny_tier_expr,
     fpl_pct_expr,
diff --git a/utils/post/apply_ny_lmi_to_master_bills.py b/utils/post/apply_ny_lmi_to_master_bills.py
@@ -41,7 +41,7 @@
 import polars as pl
 from dotenv import load_dotenv
 
-from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+from utils.file_io import get_aws_storage_options
 from utils.post.lmi_common import (
     assign_ny_tier_expr,
     fpl_pct_expr,
diff --git a/utils/post/apply_ri_lmi_discounts_to_bills.py b/utils/post/apply_ri_lmi_discounts_to_bills.py
@@ -26,7 +26,7 @@
 from cloudpathlib import S3Path
 from dotenv import load_dotenv
 
-from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+from utils.file_io import get_aws_storage_options
 from utils.post.lmi_common import (
     assign_ri_gas_tier_expr,
     assign_ri_tier_expr,
diff --git a/utils/post/build_master_bat.py b/utils/post/build_master_bat.py
@@ -58,7 +58,7 @@
 
 import polars as pl
 
-from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+from utils.file_io import get_aws_storage_options
 from utils.post.io import BLDG_ID, scan
 from utils.post.master_run12_passthrough import (
     REFERENCE_COMB_RUN_PAIR,
diff --git a/utils/post/build_master_bills.py b/utils/post/build_master_bills.py
@@ -44,7 +44,7 @@
 
 import polars as pl
 
-from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+from utils.file_io import get_aws_storage_options
 from utils.post import apply_ny_lmi_to_master_bills as ny_lmi_master_bills
 from utils.post.apply_ny_lmi_to_master_bills import apply_ny_lmi_to_master
 from utils.post.apply_ri_lmi_discounts_to_bills import apply_ri_lmi_to_master
diff --git a/utils/post/compare_master_batches.py b/utils/post/compare_master_batches.py
@@ -28,7 +28,7 @@
 import numpy as np
 import polars as pl
 
-from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+from utils.file_io import get_aws_storage_options
 from utils.post.build_master_bat import VALID_RUN_PAIRS
 from utils.post.io import BLDG_ID, path_or_s3
 
diff --git a/utils/post/validate_lmi_electric_discounts.py b/utils/post/validate_lmi_electric_discounts.py
@@ -24,7 +24,7 @@
 import polars as pl
 from dotenv import load_dotenv
 
-from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+from utils.file_io import get_aws_storage_options
 from utils.post.lmi_common import get_ny_eap_credits_df
 
 matplotlib.use("Agg")
diff --git a/utils/post/validate_lmi_gas_discounts.py b/utils/post/validate_lmi_gas_discounts.py
@@ -18,7 +18,7 @@
 import polars as pl
 from dotenv import load_dotenv
 
-from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+from utils.file_io import get_aws_storage_options
 from utils.post.lmi_common import get_ny_eap_credits_df
 
 matplotlib.use("Agg")
diff --git a/utils/post/validate_ny_mc_data.py b/utils/post/validate_ny_mc_data.py
@@ -42,7 +42,7 @@
 )
 from plotnine.composition import Compose
 
-from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+from utils.file_io import get_aws_storage_options
 from utils.data_prep.marginal_costs.supply_utils import (
     build_mc_partition_parquet_path,
     build_mc_partition_path,
diff --git a/utils/post/validate_passthrough_rr_master_vs_yaml.py b/utils/post/validate_passthrough_rr_master_vs_yaml.py
@@ -25,7 +25,7 @@
 import polars as pl
 import yaml
 
-from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+from utils.file_io import get_aws_storage_options
 from utils import get_project_root
 from utils.mid.compute_subclass_rr import parse_group_value_to_subclass
 from utils.mid.resolve_subclass_config import resolve_subclass_config
diff --git a/utils/pre/calibrate_demand_flex_elasticity.py b/utils/pre/calibrate_demand_flex_elasticity.py
@@ -37,7 +37,7 @@
 import polars as pl
 import yaml
 
-from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+from utils.file_io import get_aws_storage_options
 from utils.cairo import (
     _build_period_consumption,
     _build_period_shift_targets,
diff --git a/utils/pre/create_flat_tariffs.py b/utils/pre/create_flat_tariffs.py
@@ -39,7 +39,7 @@ def _s3_storage_options(path: str) -> dict[str, str] | None:
     if not path.startswith("s3://"):
         return None
     try:
-        from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+        from utils.file_io import get_aws_storage_options
 
         return get_aws_storage_options()
     except ImportError:
diff --git a/utils/pre/rev_requirement/analyze_rr_scatter.py b/utils/pre/rev_requirement/analyze_rr_scatter.py
@@ -25,7 +25,7 @@
     theme_minimal,
 )
 
-from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+from utils.file_io import get_aws_storage_options
 from utils.scenario_config import get_residential_customer_count_from_utility_stats
 
 UtilityData = dict[str, str | float]
diff --git a/utils/pre/rev_requirement/compute_rr.py b/utils/pre/rev_requirement/compute_rr.py
@@ -389,7 +389,7 @@ def main() -> None:
     storage_options = None
     if args.path_electric_utility_stats.startswith("s3://"):
         try:
-            from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+            from utils.file_io import get_aws_storage_options
 
             storage_options = get_aws_storage_options()
         except ImportError:
diff --git a/utils/pre/rev_requirement/estimate_psegli_rr.py b/utils/pre/rev_requirement/estimate_psegli_rr.py
@@ -273,9 +273,7 @@ def main() -> None:
     storage_options = None
     if args.path_electric_utility_stats.startswith("s3://"):
         try:
-            from data.eia.hourly_loads.eia_region_config import (
-                get_aws_storage_options,
-            )
+            from utils.file_io import get_aws_storage_options
 
             storage_options = get_aws_storage_options()
         except ImportError:
diff --git a/utils/pre/write_tariff_maps_from_scenario.py b/utils/pre/write_tariff_maps_from_scenario.py
@@ -27,7 +27,7 @@
 import polars as pl
 import yaml
 
-from data.eia.hourly_loads.eia_region_config import get_aws_storage_options
+from utils.file_io import get_aws_storage_options
 from utils.pre.electric_tariff_mapper import generate_tariff_map_from_scenario_keys
 
 _S3_OPTS = get_aws_storage_options()

Original file line number	Diff line number	Diff line change
`@@ -120,7 +120,7 @@ def main() -> None:`
`120`	`120`	`)`
`121`	`121`
`122`	`122`	`load_dotenv()`
`123`		`- from data.eia.hourly_loads.eia_region_config import get_aws_storage_options`
	`123`	`+ from utils.file_io import get_aws_storage_options`
`124`	`124`
`125`	`125`	`storage_options = get_aws_storage_options()`
`126`	`126`