Compress and ingest background files via R2D2

src/swell/tasks/get_background.py

@@ -9,6 +9,7 @@
 
 
 from swell.tasks.base.task_base import taskBase
+from swell.utilities.compress import compressed_extension, decompress_if_needed
 from swell.utilities.r2d2 import load_r2d2_credentials, get_r2d2_model_name
 
 from datetime import timedelta
@@ -189,6 +190,9 @@ def execute(self) -> None:
                 fetch_step = bkg_step
                 fetch_date = forecast_start_time.strftime('%Y-%m-%dT%H:%M:%SZ')
 
+                fetch_target = target_file
+                compress_output = False
+
                 if use_geos_cf_oper_background:
                     geos_cf_forecast_start_time = self.geos_cf_oper_forecast_start(background_time)
                     fetch_step = self.geos_cf_oper_step(
@@ -203,9 +207,16 @@ def execute(self) -> None:
                         f'date={fetch_date}, step={fetch_step}'
                     )
 
+                    # SaveBackground may have compressed this record before storing it
+                    # (compress_output). R2D2 query must match the stored extension.
+                    compress_output = self.config.compress_output(False)
+                    if compress_output:
+                        file_extension = compressed_extension(file_extension)
+                        fetch_target = target_file + '.gz'
+
                 fetch_kwargs = dict(
                     item='forecast',
-                    target_file=target_file,
+                    target_file=fetch_target,
                     model=r2d2_model,
                     experiment=background_experiment,
                     file_extension=file_extension,
@@ -218,7 +229,9 @@ def execute(self) -> None:
                     fetch_kwargs['data_store'] = r2d2_datastore
                 r2d2.fetch(**fetch_kwargs)
 
-                # Change permission
-                os.chmod(target_file, 0o644)
+                # Decompress if SaveBackground stored this record compressed, then
+                # change permission on the file that actually ends up at target_file.
+                actual_file = decompress_if_needed(fetch_target)
+                os.chmod(actual_file, 0o644)
 
 # --------------------------------------------------------------------------------------------------

src/swell/tasks/save_background.py

@@ -10,9 +10,11 @@
 from datetime import datetime as dt
 from datetime import timedelta
 import os
+import shutil
 from r2d2 import store
 
 from swell.tasks.base.task_base import taskBase
+from swell.utilities.compress import compress_file, compressed_extension
 from swell.utilities.datetime_util import datetime_formats
 from swell.utilities.r2d2 import load_r2d2_credentials
 
@@ -45,6 +47,13 @@ def execute(self) -> None:
         - ``horizontal_resolution``: R2D2 resolution string (default ``c360``)
         - ``store_as_symlink``: if ``True`` (default), register files as symlinks
           in R2D2 rather than copying them
+        - ``compress_output``: if ``True`` (default ``False``), gzip/pigz-compress
+          each file before storing to save R2D2 disk space. Incompatible with
+          ``store_as_symlink`` (the compressed copy is not the original file), so
+          ``store_as_symlink`` is forced to ``False`` when compression is enabled.
+        - ``compress_algorithm``: ``gzip`` (default) or ``pigz`` (parallel, requires
+          the ``pigz`` binary on ``PATH``)
+        - ``compress_pigz_threads``: thread count for ``pigz`` (default 4)
 
         The Cylc cycle point must be the forecast initialization time,
         e.g. ``2025-10-02T09:00:00Z``.
@@ -75,6 +84,21 @@ def execute(self) -> None:
         resolution = self.config.horizontal_resolution('c360')
         store_as_symlink = self.config.store_as_symlink(True)
 
+        compress_output = self.config.compress_output(False)
+        compress_algorithm = self.config.compress_algorithm('gzip')
+        compress_pigz_threads = self.config.compress_pigz_threads(4)
+
+        if compress_output and store_as_symlink:
+            self.logger.warning(
+                'compress_output=True and store_as_symlink=True are incompatible '
+                '(a compressed copy is not the original file). '
+                'Forcing store_as_symlink=False.'
+            )
+            store_as_symlink = False
+
+        if compress_output:
+            os.makedirs(self.cycle_dir(), 0o755, exist_ok=True)
+
         stored = 0
         skipped = 0
 
@@ -98,6 +122,32 @@ def execute(self) -> None:
 
             self.logger.info(f'  Storing step={step}: {os.path.basename(source_file)}')
 
+            store_source = source_file
+            store_extension = 'nc4'
+            staged_compressed_file = None
+
+            if compress_output:
+                # Source files live on a shared, often read-only NRT filesystem —
+                # stage a local copy in the cycle directory before compressing
+                # rather than writing next to the original.
+                staged_file = os.path.join(self.cycle_dir(), os.path.basename(source_file))
+                try:
+                    shutil.copy(source_file, staged_file)
+                    staged_compressed_file = compress_file(
+                        staged_file,
+                        algorithm=compress_algorithm,
+                        num_threads=compress_pigz_threads,
+                    )
+                except Exception as exc:
+                    self.logger.abort(
+                        f'Failed to compress background file {source_file}: {exc}')
+                finally:
+                    if os.path.exists(staged_file):
+                        os.remove(staged_file)
+
+                store_source = staged_compressed_file
+                store_extension = compressed_extension('nc4')
+
             try:
                 store(
                     model=model,
@@ -106,8 +156,8 @@ def execute(self) -> None:
                     experiment=experiment,
                     resolution=resolution,
                     date=forecast_start.strftime('%Y%m%d_%H%Mz'),
-                    source_file=source_file,
-                    file_extension='nc4',
+                    source_file=store_source,
+                    file_extension=store_extension,
                     file_type='bkg',
                     store_as_symlink=store_as_symlink,
                 )
@@ -129,6 +179,9 @@ def execute(self) -> None:
                         f'{os.path.basename(source_file)}')
                 else:
                     raise
+            finally:
+                if staged_compressed_file and os.path.exists(staged_compressed_file):
+                    os.remove(staged_compressed_file)
             stored += 1
 
         verb = 'Would store' if dry_run else 'Stored'

src/swell/tasks/task_questions.py

@@ -340,6 +340,7 @@ class TaskQuestions(QuestionContainer, Enum):
             qd.background_frequency(),
             qd.horizontal_resolution(),
             qd.marine_models(),
+            qd.compress_output(),
         ]
     )
 
@@ -956,6 +957,9 @@ class TaskQuestions(QuestionContainer, Enum):
             qd.background_experiment(),
             qd.horizontal_resolution(),
             qd.store_as_symlink(),
+            qd.compress_output(),
+            qd.compress_algorithm(),
+            qd.compress_pigz_threads(),
         ]
     )
 

src/swell/utilities/compress.py

@@ -1,4 +1,3 @@
-
 # (C) Copyright 2021- United States Government as represented by the Administrator of the
 # National Aeronautics and Space Administration. All Rights Reserved.
 #

src/swell/utilities/question_defaults.py

@@ -1023,6 +1023,48 @@ class store_as_symlink(TaskQuestion):
 
     # --------------------------------------------------------------------------------------------------
 
+    @dataclass
+    class compress_output(TaskQuestion):
+        default_value: bool = False
+        question_name: str = "compress_output"
+        ask_question: bool = False
+        models: List[str] = mutable_field([
+            "all_models"
+        ])
+        prompt: str = "Compress output files with gzip before storing in R2D2 to save disk space"
+        widget_type: WType = WType.BOOLEAN
+
+    # --------------------------------------------------------------------------------------------------
+
+    @dataclass
+    class compress_algorithm(TaskQuestion):
+        default_value: str = "gzip"
+        question_name: str = "compress_algorithm"
+        ask_question: bool = False
+        models: List[str] = mutable_field([
+            "all_models"
+        ])
+        prompt: str = ("Compression algorithm: 'gzip' (Python stdlib, no external binary) or "
+                       "'pigz' (parallel gzip, requires pigz binary on PATH)")
+        widget_type: WType = WType.STRING_DROP_LIST
+
+    # --------------------------------------------------------------------------------------------------
+
+    @dataclass
+    class compress_pigz_threads(TaskQuestion):
+        default_value: int = 4
+        question_name: str = "compress_pigz_threads"
+        ask_question: bool = False
+        models: List[str] = mutable_field([
+            "all_models"
+        ])
+        prompt: str = ("Number of parallel threads for pigz compression. "
+                       "Only used when compress_algorithm is 'pigz'. "
+                       "Should not exceed the number of CPUs allocated to the task.")
+        widget_type: WType = WType.INTEGER
+
+    # --------------------------------------------------------------------------------------------------
+
     @dataclass
     class obs_rc_path(TaskQuestion):
         default_value: str = 'GEOS_mksi/ObsClass/obsys-nccs.rc'