Optimize CDC air quality imports with sharding and scaled compute

scripts/us_cdc/environmental_health_toxicology/PM25CountyPollution.tmcf

@@ -10,7 +10,7 @@ variableMeasured: dcs:Mean_Concentration_AirPollutant_PM2.5
 Node: E:PM25CountyPollution->E2
 observationAbout: C:PM25CountyPollution->dcid
 typeOf: dcs:StatVarObservation
-observationDate: C:PPM25CountyPollution->date
+observationDate: C:PM25CountyPollution->date
 value: C:PM25CountyPollution->pm25_med_pred
 observationPeriod: "P24H"
 unit: MicrogramsPerCubicMeter

scripts/us_cdc/environmental_health_toxicology/README.md

@@ -69,19 +69,13 @@ These data were collected as part of the [CDC National Environment Public Health
 [`small_Palmer_expected.csv`](https://github.com/datacommonsorg/data/blob/master/scripts/us_cdc/environmental_health_toxicology/test_data/small_Palmer_expected.csv)
 
 #### tMCFs
-[`OzoneCensusTractPollution.tmcf`](https://github.com/datacommonsorg/data/blob/master/scripts/us_cdc/environmental_health_toxicology/OzoneCensusTractPollution.tmcf)
-
-[`OzoneCountyPollution.tmcf`](https://github.com/datacommonsorg/data/blob/master/scripts/us_cdc/environmental_health_toxicology/OzoneCountyPollution.tmcf)
-
-[`PalmerDroughtSeverityIndex.tmcf`](https://github.com/datacommonsorg/data/blob/master/scripts/us_cdc/environmental_health_toxicology/PalmerDroughtSeverityIndex.tmcf)
-
-[`PM25CensusTractPollution.tmcf`](https://github.com/datacommonsorg/data/blob/master/scripts/us_cdc/environmental_health_toxicology/PM25CensusTractPollution.tmcf)
-
-[`PM25CountyPollution.tmcf`](https://github.com/datacommonsorg/data/blob/master/scripts/us_cdc/environmental_health_toxicology/PM25CountyPollution.tmcf)
-
-[`StandardizedPrecipitationEvapotranspirationIndex.tmcf`](https://github.com/datacommonsorg/data/blob/master/scripts/us_cdc/environmental_health_toxicology/StandardizedPrecipitationEvapotranspirationIndex.tmcf)
-
-[`StandardizedPrecipitationIndex.tmcf`](https://github.com/datacommonsorg/data/blob/master/scripts/us_cdc/environmental_health_toxicology/StandardizedPrecipitationIndex.tmcf)
+* [`OzoneCountyPollution.tmcf`](https://github.com/datacommonsorg/data/blob/master/scripts/us_cdc/environmental_health_toxicology/OzoneCountyPollution.tmcf)
+* [`PM25CountyPollution.tmcf`](https://github.com/datacommonsorg/data/blob/master/scripts/us_cdc/environmental_health_toxicology/PM25CountyPollution.tmcf)
+* **Sharded PM2.5 Census Tract tMCFs:** [`PM25CensusTractPollution_part1.tmcf`](https://github.com/datacommonsorg/data/blob/master/scripts/us_cdc/environmental_health_toxicology/PM25CensusTractPollution_part1.tmcf), [`part2`](https://github.com/datacommonsorg/data/blob/master/scripts/us_cdc/environmental_health_toxicology/PM25CensusTractPollution_part2.tmcf), [`part3`](https://github.com/datacommonsorg/data/blob/master/scripts/us_cdc/environmental_health_toxicology/PM25CensusTractPollution_part3.tmcf), [`part4`](https://github.com/datacommonsorg/data/blob/master/scripts/us_cdc/environmental_health_toxicology/PM25CensusTractPollution_part4.tmcf)
+* **Sharded Ozone Census Tract tMCFs:** [`OzoneCensusTractPollution_part1.tmcf`](https://github.com/datacommonsorg/data/blob/master/scripts/us_cdc/environmental_health_toxicology/OzoneCensusTractPollution_part1.tmcf), [`part2`](https://github.com/datacommonsorg/data/blob/master/scripts/us_cdc/environmental_health_toxicology/OzoneCensusTractPollution_part2.tmcf), [`part3`](https://github.com/datacommonsorg/data/blob/master/scripts/us_cdc/environmental_health_toxicology/OzoneCensusTractPollution_part3.tmcf), [`part4`](https://github.com/datacommonsorg/data/blob/master/scripts/us_cdc/environmental_health_toxicology/OzoneCensusTractPollution_part4.tmcf)
+* [`PalmerDroughtSeverityIndex.tmcf`](https://github.com/datacommonsorg/data/blob/master/scripts/us_cdc/environmental_health_toxicology/PalmerDroughtSeverityIndex.tmcf)
+* [`StandardizedPrecipitationEvapotranspirationIndex.tmcf`](https://github.com/datacommonsorg/data/blob/master/scripts/us_cdc/environmental_health_toxicology/StandardizedPrecipitationEvapotranspirationIndex.tmcf)
+* [`StandardizedPrecipitationIndex.tmcf`](https://github.com/datacommonsorg/data/blob/master/scripts/us_cdc/environmental_health_toxicology/StandardizedPrecipitationIndex.tmcf)
 
 ### Import Procedure
 

scripts/us_cdc/environmental_health_toxicology/download_files.py

@@ -12,9 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import json, os, requests, sys
+import json
+import os
 from pathlib import Path
-from absl import app, logging, flags
+import shutil
+import sys
+
+from absl import app
+from absl import flags
+from absl import logging
+import requests
 from retry import retry
 
 _FLAGS = flags.FLAGS
@@ -36,60 +43,80 @@ def download_files(importname, configs):
     @retry(tries=3, delay=2, backoff=2)
     def download_with_retry(url, input_file_name):
         logging.info(f"Downloading file from URL: {url}")
-        response = requests.get(url)
-        response.raise_for_status()
-        if response.status_code == 200:
-            if not response.content:
-                logging.fatal(
-                    f"No data available for URL: {url}. Aborting download.")
-                return
-            filename = os.path.join(_INPUT_FILE_PATH, input_file_name)
-            with file_util.FileIO(filename, 'wb') as f:
-                f.write(response.content)
-        else:
-            logging.error(
-                f"Failed to download file from URL: {url}. Status code: {response.status_code}"
+        filename = os.path.join(_INPUT_FILE_PATH, input_file_name)
+        tmp_filename = f"{filename}.tmp"
+        try:
+            with requests.get(url, stream=True, timeout=(30, 300)) as response:
+                response.raise_for_status()
+                with open(tmp_filename, 'wb') as f:
+                    for chunk in response.iter_content(chunk_size=16 * 1024 *
+                                                       1024):
+                        if chunk:
+                            f.write(chunk)
+            if not os.path.exists(tmp_filename) or os.path.getsize(
+                    tmp_filename) <= 0:
+                raise IOError(
+                    f"Downloaded file {tmp_filename} is empty or missing.")
+            shutil.move(tmp_filename, filename)
+            logging.info(
+                f"Successfully saved {filename} ({os.path.getsize(filename)} bytes)"
             )
+        except Exception as e:
+            if os.path.exists(tmp_filename):
+                try:
+                    os.remove(tmp_filename)
+                except OSError:
+                    pass
+            raise e
+
+    @retry(tries=3, delay=2, backoff=2)
+    def get_record_count_with_retry(count_url):
+        logging.info(f"Querying record count from URL: {count_url}")
+        resp = requests.get(count_url, timeout=60)
+        resp.raise_for_status()
+        return json.loads(resp.text)[0]['COLUMN_ALIAS_GUARD__count']
 
+    url_new = None
+    import_found = False
     try:
         for config in configs:
             if config["import_name"] == importname:
+                import_found = True
                 files = config["files"]
                 for file_info in files:
                     url_new = file_info["url"]
                     logging.info(f"URL from config file {url_new}")
                     input_file_name = file_info["input_file_name"]
                     logging.info(f"Input File Name {input_file_name}")
 
-                    get_record_count = requests.get(
-                        url_new.replace('.csv', record_count_query))
-                    if get_record_count.status_code == 200:
-                        record_count = json.loads(
-                            get_record_count.text
-                        )[0]['COLUMN_ALIAS_GUARD__count']
-                        logging.info(
-                            f"Numbers of records found for the URL {url_new} is {record_count}"
-                        )
-                        url_new = f"{url_new}?$limit={record_count}&$offset=0"
-                        download_with_retry(url_new, input_file_name)
-                        logging.info(
-                            "Successfully downloaded the source data...!!!!")
-                    else:
-                        logging.error(
-                            f"Failed to download files, Status code: {get_record_count.status_code}"
-                        )
+                    count_url = url_new.replace('.csv', record_count_query)
+                    record_count = get_record_count_with_retry(count_url)
+                    logging.info(
+                        f"Numbers of records found for the URL {url_new} is {record_count}"
+                    )
+                    url_new = f"{url_new}?$limit={record_count}&$offset=0"
+                    download_with_retry(url_new, input_file_name)
+                    logging.info(
+                        "Successfully downloaded the source data...!!!!")
+        if not import_found:
+            raise ValueError(
+                f"Import name '{importname}' not found in configuration")
 
     except Exception as e:
-        logging.fatal(f"Error downloading URL {url_new} - {e}")
+        logging.fatal(f"Error downloading URL {url_new or 'unknown'} - {e}")
 
 
-def main(_):
+def main(argv):
     """Main function to download the csv files."""
+    if len(argv) < 2:
+        logging.fatal(
+            "Missing import name argument. Usage: download_files.py <import_name>"
+        )
+        return
     global _INPUT_FILE_PATH
-    _INPUT_FILE_PATH = os.path.join(_FLAGS.input_file_path)
     _INPUT_FILE_PATH = os.path.join(_MODULE_DIR, _FLAGS.input_file_path)
     Path(_INPUT_FILE_PATH).mkdir(parents=True, exist_ok=True)
-    importname = sys.argv[1]
+    importname = argv[1]
     logging.info(f'Loading config: {_FLAGS.config_file}')
     with file_util.FileIO(_FLAGS.config_file, 'r') as f:
         config = json.load(f)

scripts/us_cdc/environmental_health_toxicology/manifest.json

@@ -12,12 +12,17 @@
                 "parse_air_quality.py CDC_PM25CensusTract"
             ],
             "source_files": [
-                "input/*.gz"
+                "input_files/*",
+                "validation_config_census_tract.json"
             ],
             "resource_limits": {
-                "cpu": 8,
-                "memory": 64,
-                "disk": 100
+                "cpu": 32,
+                "memory": 512,
+                "disk": 2000
+            },
+            "validation_config_file": "validation_config_census_tract.json",
+            "config_override": {
+                "invoke_differ_tool": false
             },
             "import_inputs": [
                 {
@@ -51,12 +56,17 @@
                 "parse_air_quality.py CDC_OzoneCensusTract"
             ],
             "source_files": [
-                "input/*.gz"
+                "input_files/*",
+                "validation_config_census_tract.json"
             ],
             "resource_limits": {
-                "cpu": 8,
-                "memory": 64,
-                "disk": 100
+                "cpu": 32,
+                "memory": 512,
+                "disk": 2000
+            },
+            "validation_config_file": "validation_config_census_tract.json",
+            "config_override": {
+                "invoke_differ_tool": false
             },
             "import_inputs": [
                 {
@@ -90,19 +100,33 @@
                 "parse_air_quality.py CDC_PM25County"
             ],
             "source_files": [
-                "input_files/*"
+                "input_files/*",
+                "validation_config_county.json"
             ],
+            "resource_limits": {
+                "cpu": 32,
+                "memory": 512,
+                "disk": 500
+            },
+            "validation_config_file": "validation_config_county.json",
             "import_inputs": [
                 {
                     "template_mcf": "PM25CountyPollution.tmcf",
-                    "cleaned_csv": "output/PM25county.csv"
+                    "cleaned_csv": "output/PM25county_0.csv"
+                },
+                {
+                    "template_mcf": "PM25CountyPollution.tmcf",
+                    "cleaned_csv": "output/PM25county_1.csv"
+                },
+                {
+                    "template_mcf": "PM25CountyPollution.tmcf",
+                    "cleaned_csv": "output/PM25county_2.csv"
+                },
+                {
+                    "template_mcf": "PM25CountyPollution.tmcf",
+                    "cleaned_csv": "output/PM25county_3.csv"
                 }
             ],
-            "resource_limits": {
-                "cpu": 8,
-                "memory": 128,
-                "disk": 200
-            },
             "cron_schedule": "0 1  4 * *"
         },
         {
@@ -117,20 +141,22 @@
                 "parse_air_quality.py CDC_OzoneCounty"
             ],
             "source_files": [
-                "input_files/*"
+                "input_files/*",
+                "validation_config_county.json"
             ],
+            "resource_limits": {
+                "cpu": 32,
+                "memory": 512,
+                "disk": 500
+            },
+            "validation_config_file": "validation_config_county.json",
             "import_inputs": [
                 {
                     "template_mcf": "OzoneCountyPollution.tmcf",
                     "cleaned_csv": "output/OzoneCounty.csv"
                 }
             ],
-            "resource_limits": {
-                "cpu": 16,
-                "memory": 512,
-                "disk": 500
-            },
             "cron_schedule": "0 1  5 * *"
         }
     ]
-}
+}