Clear out task handlers for completed batches

Author: adammcdonagh

README.md

@@ -177,6 +177,7 @@ These are some environment variables that can be used to customise the behaviour
 - `OTF_LOG_JSON` - Stderr logging will be in JSON format. Set to `1` to enable
 - `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_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

@@ -114,19 +114,20 @@ When running a batch, setting `-r` (or `OTF_RUN_ID`) causes all sub-task logs to
 
 The full list of supported environment variables:
 
-| Variable                    | Description                                                                  |
-| --------------------------- | ---------------------------------------------------------------------------- |
-| `OTF_SSH_KEY`               | Default private SSH key path for all SSH connections                         |
-| `OTF_LOG_DIRECTORY`         | Base directory for log files                                                 |
-| `OTF_NO_LOG`                | Set to `1` to disable file logging                                           |
-| `OTF_LOG_JSON`              | Set to `1` for JSON-formatted stderr logging                                 |
-| `OTF_STAGING_DIR`           | Override the staging directory for file transfers (default: `/tmp`)          |
-| `OTF_RUN_ID`                | Log aggregation identifier; equivalent to `-r` flag                          |
-| `OTF_BATCH_RESUME_LOG_DATE` | Resume batch from a specific date's logs (`YYYYMMDD` format)                 |
-| `OTF_VARIABLES_FILE`        | Override the default variables file path. Comma-separate multiple files      |
-| `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_PARAMIKO_ULTRA_DEBUG`  | Set to `1` to enable ultra-verbose Paramiko SSH debug output (SFTP only)     |
+| Variable                    | Description                                                                                                                                                  |
+| --------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------ |
+| `OTF_SSH_KEY`               | Default private SSH key path for all SSH connections                                                                                                         |
+| `OTF_LOG_DIRECTORY`         | Base directory for log files                                                                                                                                 |
+| `OTF_NO_LOG`                | Set to `1` to disable file logging                                                                                                                           |
+| `OTF_LOG_JSON`              | Set to `1` for JSON-formatted stderr logging                                                                                                                 |
+| `OTF_STAGING_DIR`           | Override the staging directory for file transfers (default: `/tmp`)                                                                                          |
+| `OTF_RUN_ID`                | Log aggregation identifier; equivalent to `-r` flag                                                                                                          |
+| `OTF_BATCH_RESUME_LOG_DATE` | Resume batch from a specific date's logs (`YYYYMMDD` format)                                                                                                 |
+| `OTF_VARIABLES_FILE`        | Override the default variables file path. Comma-separate multiple files                                                                                      |
+| `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_PARAMIKO_ULTRA_DEBUG`  | Set to `1` to enable ultra-verbose Paramiko SSH debug output (SFTP only)                                                                                     |
 
 ## Overriding Variables at Runtime
 

pyproject.toml

@@ -43,6 +43,7 @@ dev = [
     "types-requests >=2.28",
     "types-paramiko >=3.0",
     "black == 26.3.1",
+    "psutil >= 5.9",
     "isort",
     "pytest",
     "bumpver",

src/opentaskpy/cli/task_run.py

@@ -30,6 +30,7 @@ def main() -> None:
                                     OTF_LOG_LEVEL - Equivalent to using -v
                                     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_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/otflogging.py

@@ -403,6 +403,7 @@ def close_log_file(logger__: logging.Logger, result: bool = False) -> None:
             log_file_name = handler.baseFilename
 
     log_handlers = []
+    loggers_to_remove = []
 
     if log_file_name:
         new_log_filename = None
@@ -413,7 +414,7 @@ def close_log_file(logger__: logging.Logger, result: bool = False) -> None:
 
         # Loop through every logger that exists and has a handler of this filename, and
         # call the close method on it. Only the last one should rename the file
-        for logger_ in list(logging.Logger.manager.loggerDict.values()):
+        for logger_name, logger_ in list(logging.Logger.manager.loggerDict.items()):
             if isinstance(logger_, logging.Logger):
                 for handler in logger_.handlers:
                     if (
@@ -422,6 +423,8 @@ def close_log_file(logger__: logging.Logger, result: bool = False) -> None:
                     ):
                         log_handlers.append(handler)
                         handler.close()
+                        loggers_to_remove.append(logger_name)
+                        break
 
         # Now everything is closed, we can rename the log file
         # If result is True, then rename the file and remove _running from the name
@@ -437,6 +440,16 @@ def close_log_file(logger__: logging.Logger, result: bool = False) -> None:
             for handler in log_handlers:
                 handler.baseFilename = new_log_filename
 
+    # Unregister all task-scoped loggers that wrote to this log file so they
+    # can be garbage collected rather than accumulating in the process.
+    for logger_name in loggers_to_remove:
+        logging.Logger.manager.loggerDict.pop(logger_name, None)
+
+    # Always unregister the logger passed in directly — this handles the case
+    # where OTF_NO_LOG is set and no file handler was attached, so the
+    # loggers_to_remove list above would be empty.
+    logging.Logger.manager.loggerDict.pop(logger__.name, None)
+
 
 def redact(log_message: str) -> str:
     """Redact any sensitive information from the log message.

src/opentaskpy/taskhandlers/batch.py

@@ -18,6 +18,7 @@
 DEFAULT_TASK_EXIT_CODE = 0
 TASK_TYPE = "B"
 BATCH_TASK_LOG_MARKER = "__OTF_BATCH_TASK_MARKER__"
+DEFAULT_BATCH_POLL_INTERVAL = 5
 
 
 class Batch(TaskHandler):
@@ -403,8 +404,11 @@ def run(self, kill_event: threading.Event | None = None) -> bool:  # noqa: C901
                 )
                 break
 
-            # Sleep 5 seconds before checking again
-            time.sleep(5)
+            # Sleep before checking again
+            poll_interval = float(
+                environ.get("OTF_BATCH_POLL_INTERVAL", DEFAULT_BATCH_POLL_INTERVAL)
+            )
+            time.sleep(poll_interval)
 
             # Check if we have been asked to kill the batch
             if kill_event and kill_event.is_set():
@@ -483,6 +487,10 @@ def task_runner(self, batch_task: dict, event: threading.Event) -> None:
                             batch_task["batch_task_spec"]["order_id"],
                             batch_task["task_id"],
                         )
+                        # Release the task handler so that remote handler objects
+                        # (boto3 clients, SSH connections, etc.) can be garbage
+                        # collected rather than accumulating for the whole batch run.
+                        batch_task["task_handler"] = None
                         break
 
                     # Check if we have been asked to kill the thread