Enhance LOCA2 data processing by introducing asset jobs and updating configurations. Refactor asset definitions to utilize new job structure, and modify sensor implementations to align with the updated asset dependencies. Update tests to reflect changes in asset return types and configurations.

Author: BenGalewsky

downscaled_climate_data/assets/loca2.py

@@ -4,29 +4,29 @@
 import requests
 import s3fs
 import xarray as xr
-from dagster import AssetExecutionContext, AssetIn, Config, EnvVar, asset
+import dagster as dg
 from dagster_aws.s3 import S3Resource
 
 import downscaled_climate_data
 
 
-class Loca2Config(Config):
+class Loca2Config(dg.Config):
     s3_key: str
     url: str = "https://cirrus.ucsd.edu/~pierce/LOCA2/CONUS_regions_split/ACCESS-CM2/cent/0p0625deg/r2i1p1f1/historical/tasmax/tasmax.ACCESS-CM2.historical.r2i1p1f1.1950-2014.LOCA_16thdeg_v20220413.cent.nc"  # NOQA E501
 
 
-@asset(
+@dg.asset(
     name="loca2_raw_netcdf",
     description="Raw LOCA2 data downloaded from the web",
     code_version=downscaled_climate_data.__version__,
     group_name="loca2"
 )
-def loca2_raw_netcdf(context: AssetExecutionContext,
+def loca2_raw_netcdf(context: dg.AssetExecutionContext,
                      config: Loca2Config,
-                     s3: S3Resource) -> dict[str, str]:
+                     s3: S3Resource) -> dg.Output:
 
-    destination_bucket = EnvVar("LOCA2_BUCKET").get_value()
-    destination_path_root = EnvVar("LOCA2_RAW_PATH_ROOT").get_value()
+    destination_bucket = dg.EnvVar("LOCA2_BUCKET").get_value()
+    destination_path_root = dg.EnvVar("LOCA2_RAW_PATH_ROOT").get_value()
 
     with requests.get(config.url, stream=True) as response:
         # Raise an exception for bad HTTP responses
@@ -44,24 +44,34 @@ def loca2_raw_netcdf(context: AssetExecutionContext,
         )
 
     context.log.info(f"Downloading data to {config.s3_key}")
-    return {
-        "bucket": destination_bucket,
-        "s3_key": config.s3_key,
-    }
+    zarr_config = ZarrConfig(
+        s3_key=config.s3_key,
+        bucket=destination_bucket,
+    )
+    return dg.MaterializeResult(
+        metadata={
+            "zarr_config": dg.MetadataValue.json(zarr_config.__dict__),
+        }
+    )
 
 
-@asset(
+class ZarrConfig(dg.Config):
+    s3_key: str
+    bucket: str
+
+
+@dg.asset(
     name="loca2_zarr",
-    ins={
-        "loca2_raw_netcdf": AssetIn()
-    },
+    deps=["loca2_raw_netcdf"],
     group_name="loca2",
     description="LOCA2 data converted to Zarr format",
     code_version=downscaled_climate_data.__version__)
-def loca2_zarr(context,
-               loca2_raw_netcdf,
-               s3: S3Resource):
-    context.log.info(f"Converting {loca2_raw_netcdf['s3_key']} to zarr")
+def loca2_zarr(context: dg.AssetExecutionContext, s3: S3Resource):
+    upstream_metadata = context.instance.get_latest_materialization_event(
+        dg.AssetKey("loca2_raw_netcdf")).asset_materialization.metadata
+
+    config = ZarrConfig(**upstream_metadata['zarr_config'].data)
+    context.log.info(f"Converting {config.s3_key} to zarr")
 
     # Initialize s3fs with the same credentials as the S3Resource
     fs = s3fs.S3FileSystem(
@@ -70,14 +80,14 @@ def loca2_zarr(context,
         endpoint_url=s3.endpoint_url
     )
 
-    raw_root = EnvVar("LOCA2_RAW_PATH_ROOT").get_value()
-    zarr_root = EnvVar("LOCA2_ZARR_PATH_ROOT").get_value()
+    raw_root = dg.EnvVar("LOCA2_RAW_PATH_ROOT").get_value()
+    zarr_root = dg.EnvVar("LOCA2_ZARR_PATH_ROOT").get_value()
     # Construct S3 paths
-    input_path = f"s3://{loca2_raw_netcdf['bucket']}/{raw_root}{loca2_raw_netcdf['s3_key']}"  # NOQA E501
+    input_path = f"s3://{config.bucket}/{raw_root}{config.s3_key}"  # NOQA E501
     context.log.info(f"Reading from {input_path}")
 
-    zarr_key = loca2_raw_netcdf['s3_key'].replace('.nc', '.zarr')
-    output_path = f"s3://{loca2_raw_netcdf['bucket']}/{zarr_root}{zarr_key}"
+    zarr_key = config.s3_key.replace('.nc', '.zarr')
+    output_path = f"s3://{config.bucket}/{zarr_root}{zarr_key}"
     context.log.info(f"Writing to {output_path}")
 
     # Read NetCDF file from S3
@@ -103,7 +113,7 @@ def loca2_zarr(context,
         ds.close()
 
 
-class ESMCatalogConfig(Config):
+class ESMCatalogConfig(dg.Config):
     data_format: str = "zarr"
     id: str = "loca2_zarr_monthly_esm_catalog"
     description: str = "LOCA2 Zarr data catalog"
@@ -143,21 +153,21 @@ def parse_key(relative_path: str, bucket: str, full_key: str) -> dict[str, str]:
     }
 
 
-@asset(
+@dg.asset(
     name="loca2_esm_catalog",
     group_name="loca2",
     description="Generate an Intake-ESM Catalog for LOCA2 datasets",
     code_version=downscaled_climate_data.__version__)
-def loca2_esm_catalog(context: AssetExecutionContext,
+def loca2_esm_catalog(context: dg.AssetExecutionContext,
                       config: ESMCatalogConfig,
                       s3: S3Resource):
 
-    bucket = EnvVar("LOCA2_BUCKET").get_value()
+    bucket = dg.EnvVar("LOCA2_BUCKET").get_value()
 
     if config.is_zarr():
-        prefix = EnvVar("LOCA2_ZARR_PATH_ROOT").get_value()
+        prefix = dg.EnvVar("LOCA2_ZARR_PATH_ROOT").get_value()
     else:
-        prefix = EnvVar("LOCA2_RAW_PATH_ROOT").get_value()
+        prefix = dg.EnvVar("LOCA2_RAW_PATH_ROOT").get_value()
 
     if config.is_monthly():
         prefix += "/monthly"

downscaled_climate_data/definitions.py

@@ -1,4 +1,4 @@
-from dagster import Definitions, EnvVar
+from dagster import Definitions, EnvVar, define_asset_job
 from dagster_aws.s3 import S3Resource
 
 from downscaled_climate_data.assets.loca2 import loca2_zarr, loca2_raw_netcdf
@@ -12,8 +12,19 @@
                                                           loca2_sensor_monthly_tasmax,
                                                           loca2_sensor_tasmin)
 
+all_assets = [loca2_raw_netcdf, loca2_zarr, loca2_esm_catalog]
+
+loca2_data_job = define_asset_job(
+    name="loca2_data_job",
+    selection=[
+        "loca2_raw_netcdf",
+        "loca2_zarr"
+    ],
+)
+
 defs = Definitions(
-    assets=[loca2_raw_netcdf, loca2_zarr, loca2_esm_catalog],
+    assets=all_assets,
+    jobs=[loca2_data_job],
     sensors=[loca2_sensor_tasmax,
              loca2_sensor_tasmin,
              loca2_sensor_pr,

downscaled_climate_data/sensors/loca2_sensor.py

@@ -12,7 +12,6 @@
     RunConfig,
 )
 
-from downscaled_climate_data.assets.loca2 import loca2_raw_netcdf, loca2_zarr
 from downscaled_climate_data.sensors.loca2_models import Loca2Models
 
 # Give ourselves 2 hours to process a single model/scenario
@@ -21,8 +20,6 @@
 # For the smaller, monthly files, we can process them more frequently
 LOCA2_MONTHLY_SENSOR_FREQUENCY = 120
 
-LOCA2_ASSETS = [loca2_raw_netcdf, loca2_zarr]
-
 
 class Loca2Datasets(ConfigurableResource):
     """
@@ -113,15 +110,15 @@ def run_request(file: dict[str, str],
     :param monthly:
     :return:
     """
+    s3_key = "/monthly" + file["s3_key"] if monthly else file["s3_key"]
     return RunRequest(
-        run_key=file["s3_key"],
+        run_key=s3_key,
         run_config=RunConfig(
             {
                 "loca2_raw_netcdf": {
                     "config": {
                         "url": file["url"],
-                        "s3_key": "/monthly" + file["s3_key"]
-                        if monthly else file["s3_key"],
+                        "s3_key": s3_key,
                     }
                 },
             }
@@ -155,7 +152,7 @@ def sensor_implementation(context, models,
 
     # Now we can launch jobs for each of the files for this model/scenario combination
     for file in dataset_resource.get_downloadable_files(
-        models, model, scenario, monthly=True
+        models, model, scenario, monthly=monthly
     ):
         context.log.info(f"Found file: {file['url']}")
         yield run_request(file, model, scenario, monthly=monthly)
@@ -165,7 +162,7 @@ def sensor_implementation(context, models,
 
 @sensor(
     name="LOCA2_Sensor_tasmax",
-    target=LOCA2_ASSETS,
+    job_name="loca2_data_job",
     minimum_interval_seconds=LOCA2_SENSOR_FREQUENCY,
     tags={
         "variable": "tasmax",
@@ -182,7 +179,7 @@ def loca2_sensor_tasmax(
 
 @sensor(
     name="LOCA2_Sensor_tasmin",
-    target=LOCA2_ASSETS,
+    job_name="loca2_data_job",
     minimum_interval_seconds=LOCA2_SENSOR_FREQUENCY,
     tags={
         "variable": "tasmin",
@@ -199,7 +196,7 @@ def loca2_sensor_tasmin(
 
 @sensor(
     name="LOCA2_Sensor_pr",
-    target=LOCA2_ASSETS,
+    job_name="loca2_data_job",
     minimum_interval_seconds=LOCA2_SENSOR_FREQUENCY,
     tags={
         "variable": "pr",
@@ -216,7 +213,7 @@ def loca2_sensor_pr(
 
 @sensor(
     name="LOCA2_Sensor_Monthly_tasmax",
-    target=LOCA2_ASSETS,
+    job_name="loca2_data_job",
     minimum_interval_seconds=LOCA2_MONTHLY_SENSOR_FREQUENCY,
     tags={
         "variable": "tasmax",
@@ -234,7 +231,7 @@ def loca2_sensor_monthly_tasmax(
 
 @sensor(
     name="LOCA2_Sensor_Monthly_tasmin",
-    target=LOCA2_ASSETS,
+    job_name="loca2_data_job",
     minimum_interval_seconds=LOCA2_MONTHLY_SENSOR_FREQUENCY,
     tags={
         "variable": "tasmin",
@@ -252,7 +249,7 @@ def loca2_sensor_monthly_tasmin(
 
 @sensor(
     name="LOCA2_Sensor_Monthly_pr",
-    target=LOCA2_ASSETS,
+    job_name="loca2_data_job",
     minimum_interval_seconds=LOCA2_MONTHLY_SENSOR_FREQUENCY,
     tags={
         "variable": "pr",

tests/assets/test_loca2_raw_netcdf.py

@@ -1,7 +1,7 @@
 import os
 from unittest.mock import patch
 
-from dagster import DagsterInstance, build_asset_context
+from dagster import DagsterInstance, MaterializeResult, build_asset_context
 
 from downscaled_climate_data.assets.loca2 import Loca2Config, loca2_raw_netcdf
 
@@ -28,8 +28,9 @@ def test_loca2_raw(mocker):
         mock_get.return_value.__enter__.return_value = mock_response
 
         results = loca2_raw_netcdf(context=ctx, config=config)
-
-        assert results == {
+        assert type(results) == MaterializeResult
+        assert 'zarr_config' in results.metadata
+        assert results.metadata['zarr_config'].data == {
             "bucket": "test_bucket",
             "s3_key": "/loca2/cent.nc"
         }

tests/assets/test_loca2_zarr.py

@@ -1,6 +1,6 @@
 import os
 from unittest.mock import patch
-from dagster import DagsterInstance, build_asset_context
+from dagster import DagsterInstance
 
 from downscaled_climate_data.assets.loca2 import loca2_zarr
 
@@ -15,10 +15,18 @@ def test_as_zarr_asset(mock_s3fs, mock_xarray, mocker):
     s3.aws_secret_access_key = "test_secret"
     s3.endpoint_url = "https://test"
 
-    ctx = build_asset_context(instance=instance,
-                              resources={
-                                  "s3": s3
-                              })
+    # Mock the upstream asset materialization
+    mock_materialization = mocker.MagicMock()
+    mock_materialization.asset_materialization.metadata = {
+        'zarr_config': mocker.MagicMock(data={
+            'bucket': 'test_bucket',
+            's3_key': '/hist/cent.nc'
+        })
+    }
+
+    # Mock the instance to return our mock materialization
+    instance.get_latest_materialization_event = mocker.MagicMock()
+    instance.get_latest_materialization_event.return_value = mock_materialization
 
     mock_s3fs.S3FileSystem = mocker.MagicMock()
     mock_fs_open = mocker.MagicMock()
@@ -31,11 +39,19 @@ def test_as_zarr_asset(mock_s3fs, mock_xarray, mocker):
 
     os.environ['LOCA2_RAW_PATH_ROOT'] = 'test'
     os.environ['LOCA2_ZARR_PATH_ROOT'] = 'test/zarr'
-
-    loca2_zarr(context=ctx, loca2_raw_netcdf={
-        "bucket": "test_bucket",
-        "s3_key": "/hist/cent.nc"
-    })
+    from dagster import materialize_to_memory
+    result = materialize_to_memory(
+        [loca2_zarr],
+        instance=instance,
+        resources={
+            "s3": s3
+        }
+    )
+    print(result)
+    # loca2_zarr(context=ctx, loca2_raw_netcdf={
+    #     "bucket": "test_bucket",
+    #     "s3_key": "/hist/cent.nc"
+    # })
 
     mock_s3fs.S3FileSystem.assert_called_with(
         key="test_key",