Counterfactual training

Project.toml

@@ -1,7 +1,7 @@
 authors = ["Patrick Altmeyer <p.altmeyer@tudelft.nl>"]
 name = "TaijaParallel"
 uuid = "bf1c2c22-5e42-4e78-8b6b-92e6c673eeb0"
-version = "1.1.3"
+version = "1.2.0"
 
 [compat]
 Aqua = "0.8"

README.md

@@ -18,7 +18,7 @@ using Pkg
 Pkg.add("TaijaParallel.jl")
 ```
 
-The development version can be installed like so:
+The development version can be installed directly from Github:
 
 ``` julia
 using Pkg

ext/MPIExt/MPIExt.jl

@@ -7,6 +7,7 @@ using MPI
 using ProgressMeter
 using TaijaBase
 using TaijaParallel
+using TaijaParallel: load_with_retry
 
 "The `MPIParallelizer` type is used to parallelize the evaluation of a function using `MPI.jl`."
 struct MPIParallelizer <: TaijaParallel.AbstractParallelizer

ext/MPIExt/evaluate.jl

@@ -50,7 +50,8 @@ function TaijaBase.parallelize(
             if parallelizer.rank == 0 && verbose
                 # Generating counterfactuals with progress bar:
                 output = []
-                @showprogress desc = "Evaluating counterfactuals using MPI ..." for x in zip(
+                @showprogress desc = "Evaluating counterfactuals using MPI ..." for x in
+                                                                                    zip(
                     eachcol(worker_chunk)...,
                 )
                     with_logger(NullLogger()) do
@@ -70,7 +71,7 @@ function TaijaBase.parallelize(
                 second_parallelizer,
                 f,
                 eachcol(worker_chunk)...;
-                verbose=verbose,
+                verbose = verbose,
                 kwargs...,
             )
         end
@@ -92,7 +93,7 @@ function TaijaBase.parallelize(
     if parallelizer.rank == 0
         outputs = []
         for i = 1:length(chunks)
-            output = Serialization.deserialize(joinpath(storage_path, "output_$i.jls"))
+            output = load_with_retry(joinpath(storage_path, "output_$i.jls"))
             push!(outputs, output)
         end
         # Collect output from all processes in rank 0:

ext/MPIExt/generate_counterfactual.jl

@@ -52,8 +52,13 @@ function TaijaBase.parallelize(
             if parallelizer.rank == 0 && verbose
                 # Generating counterfactuals with progress bar:
                 output = []
-                @showprogress desc = "Generating counterfactuals using MPI ..." for x in zip(
-                    _x, _target, fill(data, length(_generator)), _M, _generator,
+                @showprogress desc = "Generating counterfactuals using MPI ..." for x in
+                                                                                    zip(
+                    _x,
+                    _target,
+                    fill(data, length(_generator)),
+                    _M,
+                    _generator,
                 )
                     with_logger(NullLogger()) do
                         push!(output, f(x...; kwargs...))
@@ -71,7 +76,11 @@ function TaijaBase.parallelize(
             output = TaijaBase.parallelize(
                 second_parallelizer,
                 f,
-                _x, _target, data, _M, _generator;
+                _x,
+                _target,
+                data,
+                _M,
+                _generator;
                 verbose = verbose,
                 kwargs...,
             )
@@ -94,7 +103,7 @@ function TaijaBase.parallelize(
     if parallelizer.rank == 0
         outputs = []
         for i = 1:length(chunks)
-            output = Serialization.deserialize(joinpath(storage_path, "output_$i.jls"))
+            output = load_with_retry(joinpath(storage_path, "output_$i.jls"))
             push!(outputs, output)
         end
         # Collect output from all processes in rank 0:

src/CounterfactualExplanations.jl/threads/generate_counterfactual.jl

@@ -33,16 +33,19 @@ function TaijaBase.parallelize(
     if return_flattened
         ces = Vector{CounterfactualExplanations.FlattenedCE}(undef, length(args))
     else
-        ces = Vector{CounterfactualExplanations.CounterfactualExplanation}(undef, length(args))
+        ces = Vector{CounterfactualExplanations.CounterfactualExplanation}(
+            undef,
+            length(args),
+        )
     end
 
     # Verbosity setup:
     if verbose
         prog = ProgressMeter.Progress(
             length(args);
-            desc="Generating counterfactuals using multi-threading ...",
-            showspeed=true,
-            color=:green,
+            desc = "Generating counterfactuals using multi-threading ...",
+            showspeed = true,
+            color = :green,
         )
     end
 

src/utils.jl

@@ -1,3 +1,5 @@
+using Serialization
+
 """
     chunk_obs(obs::AbstractVector, n_each::Integer, n_groups::Integer)
 
@@ -38,3 +40,23 @@ function split_obs(obs::AbstractVector, n::Integer)
     N_counts = split_count(N, n)
     return split_by_counts(obs, N_counts)
 end
+
+"""
+    load_with_retry(filepath; max_attempts=5, delay=1.0)
+
+Load a file using Serialization.deserialize, retrying up to `max_attempts` times with exponential backoff.
+"""
+function load_with_retry(filepath; max_attempts = 5, delay = 1.0)
+    for attempt = 1:max_attempts
+        try
+            return Serialization.deserialize(filepath)
+        catch e
+            if isa(e, EOFError) && attempt < max_attempts
+                sleep(delay * attempt)  # Exponential backoff
+                continue
+            end
+            rethrow(e)  # Re-throw if it's not an EOFError or we're out of attempts
+        end
+    end
+    error("Failed to load $filepath after $max_attempts attempts")
+end

test/CounterfactualExplanations.jl/comparison.jl

@@ -13,7 +13,7 @@ data = TaijaData.load_linearly_separable()
 counterfactual_data = CounterfactualData(data[1], data[2])
 
 M = fit_model(counterfactual_data, :MLP)
-conv = DecisionThresholdConvergence(decision_threshold=0.95)
+conv = DecisionThresholdConvergence(decision_threshold = 0.95)
 generator = GenericGenerator()
 factual = 1
 target = 2
@@ -23,21 +23,45 @@ xs = select_factual(counterfactual_data, chosen)
 # No parallelizer
 parallelizer = nothing
 ces = @with_parallelizer parallelizer begin
-    generate_counterfactual(xs, target, counterfactual_data, M, generator; convergence=conv, initialization=:identity)
+    generate_counterfactual(
+        xs,
+        target,
+        counterfactual_data,
+        M,
+        generator;
+        convergence = conv,
+        initialization = :identity,
+    )
 end
 
 # Threads
 parallelizer = ThreadsParallelizer()
 ces_threads = @with_parallelizer parallelizer begin
-    generate_counterfactual(xs, target, counterfactual_data, M, generator; convergence=conv, initialization=:identity)
+    generate_counterfactual(
+        xs,
+        target,
+        counterfactual_data,
+        M,
+        generator;
+        convergence = conv,
+        initialization = :identity,
+    )
 end
 
 # MPI
 using MPI
 MPI.Init()
 parallelizer = TaijaParallel.MPIParallelizer(MPI.COMM_WORLD)
 ces_mpi = @with_parallelizer parallelizer begin
-    generate_counterfactual(xs, target, counterfactual_data, M, generator; convergence=conv, initialization=:identity)
+    generate_counterfactual(
+        xs,
+        target,
+        counterfactual_data,
+        M,
+        generator;
+        convergence = conv,
+        initialization = :identity,
+    )
 end
 
 @test all(counterfactual.(ces) .== counterfactual.(ces_threads))

test/CounterfactualExplanations.jl/mpi.jl

@@ -30,13 +30,13 @@ target = fill(2, length(xs))
 ces = TaijaParallel.parallelize(
     parallelizer,
     CounterfactualExplanations.generate_counterfactual,
-    xs, 
-    target, 
-    counterfactual_data, 
-    M, 
-    generator; 
-    convergence=conv, 
-    initialization=:identity,
+    xs,
+    target,
+    counterfactual_data,
+    M,
+    generator;
+    convergence = conv,
+    initialization = :identity,
 )
 
 nsteps = (ce -> total_steps(ce)).(ces)
@@ -47,7 +47,7 @@ end
 
 # Benchmark CE with MPI
 with_logger(NullLogger()) do
-    bmk = benchmark(counterfactual_data; parallelizer=parallelizer)
+    bmk = benchmark(counterfactual_data; parallelizer = parallelizer)
 end
 MPI.Finalize()
 @test MPI.Finalized()

test/CounterfactualExplanations.jl/threads.jl

@@ -32,6 +32,10 @@ evals = @with_parallelizer parallelizer begin
     evaluate(ces)
 end
 
-bmk = benchmark(counterfactual_data; convergence=:generator_conditions, parallelizer=parallelizer)
+bmk = benchmark(
+    counterfactual_data;
+    convergence = :generator_conditions,
+    parallelizer = parallelizer,
+)
 
 @test true