Add memory usage logging env var

Author: adammcdonagh

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+# v26.18.1
+
+- Fix memory leak in batches caused by now tidying up the remote handler objects after a task has completed.
+- Fixed potential but very small memory leak in logger where loggers were not being removed from the logger dictionary after they were closed.
+- Add `OTF_LOG_MEMORY_USAGE` environment variable to allow logging of RSS memory usage at each batch poll interval using the batch logger. This is useful for diagnosing memory growth in long-running batches. Added optional dev dependency `psutil` to allow this. Needs to be installed manually, otherwise log will just output `-1.0` for RSS memory usage.
+
 # v26.18.0
 
 - Disable logging of initialisation events by default, can be enabled by setting env var `OTG_LOG_INIT_EVENTS=1`

README.md

@@ -178,6 +178,7 @@ These are some environment variables that can be used to customise the behaviour
 - `OTF_LOG_DIRECTORY` - Path under which log files are written
 - `OTF_NO_THREAD_SLEEP` - Disable the 1-second sleep between batch task thread creation. This sleep exists to prevent race conditions with concurrent protocol imports. Only disable this if you understand the implications. Set to `1` to disable.
 - `OTF_BATCH_POLL_INTERVAL` - Interval in seconds between batch status-check loop iterations (default: `5`). Reducing this makes batches respond faster when tasks complete, at the cost of higher CPU polling overhead. **Intended for testing only — do not change in production unless you fully understand the trade-offs.**
+- `OTF_LOG_MEMORY_USAGE` - Set to `1` to log the RSS memory usage of the OTF process at each batch poll interval using the batch logger. Useful for diagnosing memory growth in long-running batches. Requires `psutil` to be installed separately.
 - `OTF_RUN_ID` - (meant for internal use) An aggregator for log files. When set, all log files for a run will go under this sub directory. E.g. running a batch, all execution and transfer logs will be dropped into this sub directory, rather than a directory for each task name. This is equivalent to using `-r` or `--runId` command line arguments, which is generally preferred.
 - `OTF_SSH_KEY` - The private SSH key to use by default for all SSH connections. This is essential when using a basic docker container to trigger OTF. If not specified, it will default to use any private SSH keys available to the user executing the application.
 - `OTF_STAGING_DIR` - Staging base directory to place files before they're dropped into their final location. Default is `/tmp`

docs/usage.md

@@ -127,6 +127,7 @@ The full list of supported environment variables:
 | `OTF_LAZY_LOAD_VARIABLES`   | Set to `1` to only resolve variables that are referenced by the current task                                                                                 |
 | `OTF_NO_THREAD_SLEEP`       | Set to `1` to disable the 1-second sleep between batch task thread creation                                                                                  |
 | `OTF_BATCH_POLL_INTERVAL`   | Interval in seconds between batch status-check iterations (default: `5`). **Testing only** — do not lower in production without understanding the trade-offs |
+| `OTF_LOG_MEMORY_USAGE`      | Set to `1` to log RSS memory usage at each batch poll interval via the batch logger. Requires `psutil` to be installed                                       |
 | `OTF_PARAMIKO_ULTRA_DEBUG`  | Set to `1` to enable ultra-verbose Paramiko SSH debug output (SFTP only)                                                                                     |
 
 ## Overriding Variables at Runtime

src/opentaskpy/cli/task_run.py

@@ -31,6 +31,7 @@ def main() -> None:
                                     OTF_LOG_INIT_EVENTS - Enable logging of log initialisation events. Set to 1 to enable
                                     OTF_NO_THREAD_SLEEP - Disable the 1-second sleep between batch task thread creation. Set to 1 to disable
                                     OTF_BATCH_POLL_INTERVAL - Interval in seconds between batch status-check iterations (default: 5). Intended for testing only; do not lower in production without understanding the trade-offs
+                                    OTF_LOG_MEMORY_USAGE - Set to 1 to log RSS memory usage at each batch poll interval via the batch logger. Requires psutil
                                     OTF_SSH_KEY - Specify a particular SSH key to use for SSH/SFTP related transfers
                                     OTF_STAGING_DIR - Staging base directory to place files before final location. Default is /tmp
                                     OTF_BATCH_RESUME_LOG_DATE - Resume batch runs from a specific date in YYYYMMDD format

src/opentaskpy/taskhandlers/batch.py

@@ -1,5 +1,6 @@
 """Batch task handler class."""
 
+import os
 import threading
 import time
 from concurrent.futures import ThreadPoolExecutor, wait
@@ -21,6 +22,16 @@
 DEFAULT_BATCH_POLL_INTERVAL = 5
 
 
+def _rss_mb() -> float:
+    """Return the RSS of the current process in MB, or -1 if psutil is unavailable."""
+    try:
+        import psutil  # pylint: disable=import-outside-toplevel
+
+        return float(psutil.Process(os.getpid()).memory_info().rss / (1024 * 1024))
+    except Exception:  # pylint: disable=broad-except
+        return -1.0
+
+
 class Batch(TaskHandler):
     """Batch task handler class."""
 
@@ -410,6 +421,10 @@ def run(self, kill_event: threading.Event | None = None) -> bool:  # noqa: C901
             )
             time.sleep(poll_interval)
 
+            # Log memory usage if requested
+            if environ.get("OTF_LOG_MEMORY_USAGE"):
+                self.logger.info(f"[memory] RSS: {_rss_mb():.1f} MB")
+
             # Check if we have been asked to kill the batch
             if kill_event and kill_event.is_set():
                 self.logger.info("Kill event received, stopping threads")

tests/test_taskhandler_batch.py

@@ -676,3 +676,58 @@ def test_protocol_cache_reuse(env_vars, root_dir, clear_logs):
 
     # The second batch should use the cached protocol class
     # This can be verified through logs, but we're mainly ensuring no errors occur
+
+
+def test_batch_log_memory_usage(
+    env_vars, root_dir, clear_logs, no_thread_sleep, caplog
+):
+    """Test that OTF_LOG_MEMORY_USAGE causes RSS memory to be logged each poll cycle."""
+    import logging
+
+    old_value = os.environ.get("OTF_LOG_MEMORY_USAGE")
+    os.environ["OTF_LOG_MEMORY_USAGE"] = "1"
+    # Use a fast poll interval so the test doesn't wait on the default 5s
+    old_poll = os.environ.get("OTF_BATCH_POLL_INTERVAL")
+    os.environ["OTF_BATCH_POLL_INTERVAL"] = "0.1"
+    try:
+        mem_batch = {
+            "type": "batch",
+            "tasks": [{"order_id": i, "task_id": "df-local"} for i in range(1, 4)],
+        }
+        config_loader = ConfigLoader("test/cfg")
+        batch_obj = batch.Batch(None, f"mem-log-{RANDOM}", mem_batch, config_loader)
+
+        with caplog.at_level(logging.INFO):
+            result = batch_obj.run()
+
+        assert result, "Batch should complete successfully"
+
+        memory_log_lines = [
+            r.message for r in caplog.records if "[memory] RSS:" in r.message
+        ]
+        assert len(memory_log_lines) > 0, (
+            "Expected at least one '[memory] RSS:' log entry when OTF_LOG_MEMORY_USAGE=1, "
+            f"but none were found. All log messages: {[r.message for r in caplog.records]}"
+        )
+        # Each entry should contain a positive numeric MB value — a value of -1.0
+        # means psutil was not importable, which is a missing-dependency error.
+        import re
+
+        for line in memory_log_lines:
+            assert re.search(
+                r"\[memory\] RSS: \d+\.\d+ MB", line
+            ), f"Memory log line has unexpected format: {line!r}"
+            rss_value = float(re.search(r"\[memory\] RSS: ([\d.]+) MB", line).group(1))
+            assert rss_value > 0, (
+                f"RSS value is {rss_value} MB — a value of -1.0 means psutil is not "
+                "installed. Add 'psutil' to the project dependencies."
+            )
+    finally:
+        if old_value is None:
+            os.environ.pop("OTF_LOG_MEMORY_USAGE", None)
+        else:
+            os.environ["OTF_LOG_MEMORY_USAGE"] = old_value
+        if old_poll is None:
+            os.environ.pop("OTF_BATCH_POLL_INTERVAL", None)
+        else:
+            os.environ["OTF_BATCH_POLL_INTERVAL"] = old_poll