Enhancement: Implement customizable post-task cleanup policies (#1611)

Bourreau/app/models/bourreau_worker.rb

@@ -503,6 +503,7 @@ def process_task(task) # when entering this method, +task+ is a partial object,
     end # case 'status' is 'New', 'Data Ready', 'Recover*' and 'Restart*'
 
 
+    enqueue_post_task_cleanup_policies(task)
 
     #####################################################################
     # Task notification section
@@ -604,4 +605,38 @@ def local_cpu_quota_exceeded?(user_id) # will also create a Message if none exis
     return true # means quota is exceeded
   end
 
+  # Enqueues the post-task cleanup policies for a task that has just finished
+  def enqueue_post_task_cleanup_policies(task)
+    return unless task.status.in?(CbrainTask::FINAL_STATUS)
+
+    if task.should_cleanup_component?(:workdir)
+     BackgroundActivity::RemoveTaskWorkdir.new(
+        :user_id            => task.user_id,
+        :remote_resource_id => task.bourreau_id,
+        :status             => 'InProgress',
+        :items              => [ task.id ]
+      ).save
+    end
+
+    if task.should_cleanup_component?(:inputs)
+      BackgroundActivity::RemoveTaskInputs.new(
+        :user_id            => task.user_id,
+        :remote_resource_id => task.bourreau_id,
+        :status             => 'InProgress',
+        :items              => [ task.id ]
+      ).save
+    end
+
+    if task.should_cleanup_component?(:outputs)
+      BackgroundActivity::RemoveTaskOutputs.new(
+        :user_id            => task.user_id,
+        :remote_resource_id => task.bourreau_id,
+        :status             => 'InProgress',
+        :items              => [ task.id ]
+      ).save
+    end
+  rescue => e
+    Rails.logger.error "Post-task cleanup policies hook crashed for task ##{task.id}: #{e.message}"
+  end
+
 end

BrainPortal/app/controllers/tasks_controller.rb

@@ -1351,6 +1351,7 @@ def task_params #:nodoc:
       :user_id, :group_id, :description,
       :bourreau_id, :tool_config_id,
       :batch_id,
+      :success_cleanup_policy, :failure_cleanup_policy,
       :results_data_provider_id, :params => {}
     )
     # There are way too many 'params' in the next bit of code. Two different

BrainPortal/app/models/background_activity/remove_task_inputs.rb

@@ -0,0 +1,48 @@
+#
+# CBRAIN Project
+#
+# Copyright (C) 2008-2026
+# The Royal Institution for the Advancement of Learning
+# McGill University
+#
+
+# Removes the cached inputs of a CBRAIN task safely from the local cache.
+# Must be run on a Bourreau only.
+class BackgroundActivity::RemoveTaskInputs < BackgroundActivity::TerminateTask
+
+  Revision_info = CbrainFileRevision[__FILE__] #:nodoc:
+
+  def process(item)
+    super(item)
+
+    cbrain_task = CbrainTask.where(:bourreau_id => CBRAIN::SelfRemoteResourceId).find_by(id: item)
+    return [false, "Task not found"] unless cbrain_task
+
+    input_userfile_ids = cbrain_task.params[:interface_userfile_ids] || []
+
+    input_userfile_ids.each do |userfile_id|
+      userfile = Userfile.find_by(id: userfile_id)
+      next unless userfile
+
+      # Safety check: skip if another active task relies on this exact input
+      file_in_use = CbrainTask.active
+                              .where.not(id: cbrain_task.id)
+                              .any? { |t| t.params[:interface_userfile_ids]&.include?(userfile_id) }
+
+      if file_in_use
+        self.addlog("Skipped cache erase for '#{userfile.name}' (ID: #{userfile_id}) - asset is currently in use.")
+        next
+      end
+
+      userfile.cache_erase
+      self.addlog("Deleted cache of input file '#{userfile.name}' (ID: #{userfile_id}) via post-task cleanup policy.")
+    end
+
+    [true, nil]
+  end
+
+  def prepare_dynamic_items
+    populate_items_from_task_custom_filter
+  end
+
+end

BrainPortal/app/models/background_activity/remove_task_outputs.rb

@@ -0,0 +1,52 @@
+#
+# CBRAIN Project
+#
+# Copyright (C) 2008-2026
+# The Royal Institution for the Advancement of Learning
+# McGill University
+#
+
+# Removes the cached outputs of a CBRAIN task safely from the local cache.
+# Must be run on a Bourreau only.
+class BackgroundActivity::RemoveTaskOutputs < BackgroundActivity::TerminateTask
+
+  Revision_info = CbrainFileRevision[__FILE__] #:nodoc:
+
+  def process(item)
+    super(item)
+
+    cbrain_task = CbrainTask.where(:bourreau_id => CBRAIN::SelfRemoteResourceId).find_by(id: item)
+    return [false, "Task not found"] unless cbrain_task
+
+    output_userfiles = Userfile.where(task_id: cbrain_task.id).to_a
+
+    if output_userfiles.empty?
+      self.addlog("No output userfiles registered to erase for task ##{item}.")
+      return [true, nil]
+    end
+
+    output_userfiles.each do |userfile|
+      userfile_id = userfile.id
+
+      # Safety check: ensure no downstream task is reading this output file
+      output_in_use = CbrainTask.active
+                                .where.not(id: cbrain_task.id)
+                                .any? { |t| t.params[:interface_userfile_ids]&.include?(userfile_id) }
+
+      if output_in_use
+        self.addlog("Skipped output cache erase for '#{userfile.name}' (ID: #{userfile_id}) - asset is in use downstream.")
+        next
+      end
+
+      userfile.cache_erase
+      self.addlog("Deleted cache of output file '#{userfile.name}' (ID: #{userfile_id}) via post-task cleanup policy.")
+    end
+
+    [true, nil]
+  end
+
+  def prepare_dynamic_items
+    populate_items_from_task_custom_filter
+  end
+
+end

BrainPortal/app/models/cbrain_task.rb

@@ -37,6 +37,7 @@ class CbrainTask < ApplicationRecord
   include ResourceAccess
 
   before_validation     :set_group
+  before_validation     :set_cleanup_policy_defaults, on: :create
   after_save            :after_save_set_batch_id
   after_destroy         :remove_workdir_archive
 
@@ -264,6 +265,24 @@ def self.properties
     }
   end
 
+  def should_cleanup_component?(component)
+    if self.status == 'Completed'
+     case success_cleanup_policy
+      when 'erase_workdir'             then component == :workdir
+     when 'erase_workdir_and_outputs'  then [:workdir, :outputs].include?(component)
+     when 'erase_all'                 then [:workdir, :inputs, :outputs].include?(component)
+     when 'erase_caches'              then [:inputs, :outputs].include?(component)
+     else false
+     end
+   elsif self.status == 'Failed' || self.status.in?(CbrainTask::FAILED_STATUS)
+     case failure_cleanup_policy
+      when 'erase_workdir'             then component == :workdir
+      else false
+     end
+   else
+     false
+    end
+  end
 
   ##################################################################
   # Utility Methods
@@ -1006,6 +1025,11 @@ def set_group #:nodoc:
     end
   end
 
+  def set_cleanup_policy_defaults
+    self.success_cleanup_policy ||= 'erase_workdir_and_outputs'
+    self.failure_cleanup_policy ||= 'keep_all'
+  end
+
   def remove_workdir_archive #:nodoc:
     archive = self.workdir_archive
     return true unless archive

BrainPortal/app/views/tasks/_control.html.erb

@@ -73,5 +73,26 @@
     </td>
   </tr>
 
+  <tr>
+    <td>
+      <strong>On Task Success (Completed):</strong><br>
+      <%= select_tag "cbrain_task[success_cleanup_policy]", options_for_select([
+        ['Keep everything (Current behavior)', 'keep_all'],
+        ['Erase work directory, keep cache', 'erase_workdir'],
+        ['Erase work directory & output cache (Recommended Default)', 'erase_workdir_and_outputs'],
+        ['Erase everything (Workdir, input & output cache)', 'erase_all'],
+        ['Keep work directory, erase input & output caches', 'erase_caches']
+      ], @task.success_cleanup_policy) %>
+    </td>
+
+    <td>
+      <strong>On Task Failure (Failed):</strong><br>
+      <%= select_tag "cbrain_task[failure_cleanup_policy]", options_for_select([
+        ['Keep everything (Current behavior)', 'keep_all'],
+        ['Erase work directory, keep input cache', 'erase_workdir']
+      ], @task.failure_cleanup_policy) %>
+    </td>
+  </tr>
+
 </table>
 

BrainPortal/db/migrate/20260629232300_add_cleanup_policies_to_cbrain_tasks.rb

@@ -0,0 +1,6 @@
+class AddCleanupPoliciesToCbrainTasks < ActiveRecord::Migration[5.0]
+  def change
+    add_column :cbrain_tasks, :success_cleanup_policy, :string, default: 'keep_all'
+    add_column :cbrain_tasks, :failure_cleanup_policy, :string, default: 'keep_all'
+  end
+end