feat: federated learning demo on the actor framework
- supervisor: ignore stale death messages (OneForAll restart storm), drop normally-exited children instead of respawning, add a restart-rate limit - register actors before their thread starts so UUID dispatch and getSelf cannot race the spawn - killActor now kills the thread so links see Killed - fulfill connection promises with the pool winner on simultaneous connects - clamp the PaySim test split on small datasets and split shared transaction types across trainers in non-IID partitioning
This commit is contained in:
@@ -1 +1,3 @@
|
||||
dist-newstyle
|
||||
demo-out/
|
||||
fl_model_*.bin
|
||||
|
||||
+38
-1
@@ -71,6 +71,19 @@ library
|
||||
Control.Actor.Supervision
|
||||
Control.Actor.Network
|
||||
Control.Actor.Registry
|
||||
FL.Types
|
||||
FL.CRDT
|
||||
FL.NN
|
||||
FL.Retry
|
||||
FL.Data.PaySim
|
||||
FL.Actors.Logger
|
||||
FL.Actors.Evaluator
|
||||
FL.Actors.Trainer
|
||||
FL.Actors.Aggregator
|
||||
FL.Actors.Coordinator
|
||||
FL.Actors.PeerCoordinator
|
||||
FL.Provider
|
||||
FL.P2P
|
||||
|
||||
-- Modules included in this library but not exported.
|
||||
-- other-modules:
|
||||
@@ -79,10 +92,34 @@ library
|
||||
-- other-extensions:
|
||||
|
||||
-- Other library packages from which modules are imported.
|
||||
build-depends: base ^>=4.18.3.0, stm, uuid, bytestring ^>=0.12.0.0, containers ^>=0.8, mtl ^>=2.3.0, binary ^>=0.8.9, network ^>=3.2
|
||||
build-depends:
|
||||
base ^>=4.18.3.0,
|
||||
stm,
|
||||
uuid,
|
||||
bytestring ^>=0.12.0.0,
|
||||
containers ^>=0.8,
|
||||
mtl ^>=2.3.0,
|
||||
binary ^>=0.8.9,
|
||||
network ^>=3.2,
|
||||
hmatrix,
|
||||
cassava,
|
||||
vector,
|
||||
random,
|
||||
text,
|
||||
time
|
||||
|
||||
-- Directories containing source files.
|
||||
hs-source-dirs: src
|
||||
|
||||
-- Base language which the package is written in.
|
||||
default-language: GHC2021
|
||||
|
||||
executable fl-actors
|
||||
import: warnings
|
||||
main-is: Main.hs
|
||||
hs-source-dirs: app
|
||||
build-depends:
|
||||
base ^>=4.18.3.0,
|
||||
actors,
|
||||
optparse-applicative
|
||||
default-language: GHC2021
|
||||
|
||||
+150
@@ -0,0 +1,150 @@
|
||||
module Main where
|
||||
|
||||
import Control.Actor.Types (NodeAddr (..))
|
||||
import FL.P2P (runPeer)
|
||||
import FL.Provider (runCoordinator, runTrainer)
|
||||
import FL.Types
|
||||
import Options.Applicative
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- CLI
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
data Mode = Provider | P2P deriving (Show, Eq)
|
||||
|
||||
data Role = Coordinator | Trainer | Peer deriving (Show, Eq)
|
||||
|
||||
data CLIArgs = CLIArgs
|
||||
{ argMode :: Mode
|
||||
, argRole :: Role
|
||||
, argPort :: Int
|
||||
, argBasePort :: Int
|
||||
, argRounds :: Int
|
||||
, argEpochs :: Int
|
||||
, argLR :: Double
|
||||
, argBatch :: Int
|
||||
, argNumTrainers :: Int
|
||||
, argTrainerId :: Int
|
||||
, argCoordHost :: String
|
||||
, argCoordPort :: Int
|
||||
, argKnownPeers :: [String]
|
||||
, argDataPath :: Maybe FilePath
|
||||
, argPartition :: PartitionMode
|
||||
, argNorm :: NormMode
|
||||
}
|
||||
|
||||
parseMode :: String -> Either String Mode
|
||||
parseMode "provider" = Right Provider
|
||||
parseMode "p2p" = Right P2P
|
||||
parseMode s = Left ("unknown mode: " <> s)
|
||||
|
||||
parsePartition :: String -> Either String PartitionMode
|
||||
parsePartition "iid" = Right IID
|
||||
parsePartition "noniid" = Right NonIID
|
||||
parsePartition s = Left ("unknown partition: " <> s)
|
||||
|
||||
parseNorm :: String -> Either String NormMode
|
||||
parseNorm "minmax" = Right MinMax
|
||||
parseNorm "zscore" = Right ZScore
|
||||
parseNorm s = Left ("unknown norm: " <> s)
|
||||
|
||||
parseRole :: String -> Either String Role
|
||||
parseRole "coordinator" = Right Coordinator
|
||||
parseRole "trainer" = Right Trainer
|
||||
parseRole "peer" = Right Peer
|
||||
parseRole s = Left ("unknown role: " <> s)
|
||||
|
||||
parsePeer :: String -> NodeAddr
|
||||
parsePeer s =
|
||||
let (host, rest) = break (== ':') s
|
||||
in NodeAddr host (read (drop 1 rest))
|
||||
|
||||
argsParser :: Parser CLIArgs
|
||||
argsParser = CLIArgs
|
||||
<$> option (eitherReader parseMode)
|
||||
(long "mode" <> short 'm' <> metavar "provider|p2p"
|
||||
<> help "Topology mode" <> value Provider <> showDefaultWith show)
|
||||
<*> option (eitherReader parseRole)
|
||||
(long "role" <> short 'r' <> metavar "coordinator|trainer|peer"
|
||||
<> help "Node role" <> value Coordinator <> showDefaultWith show)
|
||||
<*> option auto
|
||||
(long "port" <> metavar "PORT"
|
||||
<> help "This node's port" <> value 9100 <> showDefault)
|
||||
<*> option auto
|
||||
(long "base-port" <> metavar "PORT"
|
||||
<> help "Base port (coordinator uses this)" <> value 9100 <> showDefault)
|
||||
<*> option auto
|
||||
(long "rounds" <> metavar "N"
|
||||
<> help "Number of federated rounds" <> value 10 <> showDefault)
|
||||
<*> option auto
|
||||
(long "epochs" <> short 'e' <> metavar "N"
|
||||
<> help "Local training epochs per round" <> value 5 <> showDefault)
|
||||
<*> option auto
|
||||
(long "lr" <> metavar "LR"
|
||||
<> help "Learning rate" <> value 0.01 <> showDefault)
|
||||
<*> option auto
|
||||
(long "batch" <> metavar "N"
|
||||
<> help "Mini-batch size" <> value 32 <> showDefault)
|
||||
<*> option auto
|
||||
(long "trainers" <> short 't' <> metavar "N"
|
||||
<> help "Expected number of trainers/peers" <> value 3 <> showDefault)
|
||||
<*> option auto
|
||||
(long "id" <> metavar "N"
|
||||
<> help "Trainer/peer index (0-based)" <> value 0 <> showDefault)
|
||||
<*> strOption
|
||||
(long "coord-host" <> metavar "HOST"
|
||||
<> help "Coordinator hostname (trainer mode)" <> value "localhost")
|
||||
<*> option auto
|
||||
(long "coord-port" <> metavar "PORT"
|
||||
<> help "Coordinator port (trainer mode)" <> value 9100 <> showDefault)
|
||||
<*> many (strOption
|
||||
(long "peer" <> metavar "HOST:PORT"
|
||||
<> help "Known peer address (p2p mode, repeat for each peer)"))
|
||||
<*> optional (strOption
|
||||
(long "data" <> short 'd' <> metavar "PATH"
|
||||
<> help "Path to PaySim CSV (uses synthetic data if absent)"))
|
||||
<*> option (eitherReader parsePartition)
|
||||
(long "partition" <> metavar "iid|noniid"
|
||||
<> help "Data partition mode" <> value IID <> showDefaultWith show)
|
||||
<*> option (eitherReader parseNorm)
|
||||
(long "norm" <> metavar "minmax|zscore"
|
||||
<> help "Feature normalization mode" <> value MinMax <> showDefaultWith show)
|
||||
|
||||
main :: IO ()
|
||||
main = do
|
||||
args <- execParser $
|
||||
info (argsParser <**> helper)
|
||||
( fullDesc
|
||||
<> progDesc "Federated Learning on the Haskell Actor Framework"
|
||||
<> header "fl-actors — FedAvg fraud detection"
|
||||
)
|
||||
|
||||
let cfg = FLConfig
|
||||
{ cfgRounds = argRounds args
|
||||
, cfgEpochs = argEpochs args
|
||||
, cfgLearningRate = argLR args
|
||||
, cfgBatchSize = argBatch args
|
||||
, cfgNumTrainers = argNumTrainers args
|
||||
, cfgDataPath = argDataPath args
|
||||
, cfgPartition = argPartition args
|
||||
, cfgNorm = argNorm args
|
||||
, cfgBasePort = argBasePort args
|
||||
}
|
||||
|
||||
case (argMode args, argRole args) of
|
||||
(Provider, Coordinator) ->
|
||||
runCoordinator cfg { cfgBasePort = argPort args }
|
||||
|
||||
(Provider, Trainer) -> do
|
||||
let coordAddr = NodeAddr (argCoordHost args) (fromIntegral (argCoordPort args))
|
||||
runTrainer cfg (argTrainerId args) coordAddr
|
||||
|
||||
(P2P, Peer) -> do
|
||||
let peers = map parsePeer (argKnownPeers args)
|
||||
cfg' = cfg { cfgBasePort = argPort args }
|
||||
runPeer cfg' (argTrainerId args) peers
|
||||
|
||||
(mode, role) ->
|
||||
putStrLn $ "Invalid combination: mode=" <> show mode <> " role=" <> show role
|
||||
<> "\n provider mode uses: --role=coordinator or --role=trainer"
|
||||
<> "\n p2p mode uses: --role=peer"
|
||||
Executable
+85
@@ -0,0 +1,85 @@
|
||||
#!/usr/bin/env bash
|
||||
#
|
||||
# Demo: federated learning in P2P mode (full mesh, no central coordinator).
|
||||
# Starts N peer nodes as separate processes; each trains locally, gossips
|
||||
# updates via PeerSync (CRDT-gated rounds), and aggregates independently.
|
||||
# At the end every peer must hold the SAME model — the script verifies the
|
||||
# saved model files are byte-identical.
|
||||
#
|
||||
# Usage:
|
||||
# scripts/demo-p2p.sh [ROUNDS] [PEERS] [EPOCHS] [/path/to/paysim.csv]
|
||||
#
|
||||
# Without a CSV path, nodes fall back to synthetic data (fast; the metrics are
|
||||
# meaningless, but the gossip protocol and aggregation are fully exercised).
|
||||
set -euo pipefail
|
||||
|
||||
ROUNDS="${1:-3}"
|
||||
PEERS="${2:-3}"
|
||||
EPOCHS="${3:-2}"
|
||||
DATA="${4:-}"
|
||||
|
||||
BASE_PORT=9600
|
||||
OUT="demo-out/p2p"
|
||||
|
||||
cd "$(dirname "$0")/.."
|
||||
echo "==> Building..."
|
||||
cabal build exe:fl-actors >/dev/null
|
||||
BIN="$(cabal list-bin fl-actors)"
|
||||
|
||||
mkdir -p "$OUT"
|
||||
rm -f "$OUT"/*.log fl_model_peer-*.bin
|
||||
|
||||
DATA_ARGS=()
|
||||
[ -n "$DATA" ] && DATA_ARGS=(--data "$DATA")
|
||||
|
||||
PIDS=()
|
||||
cleanup() { kill "${PIDS[@]}" 2>/dev/null || true; wait 2>/dev/null || true; }
|
||||
trap cleanup EXIT INT TERM
|
||||
|
||||
for i in $(seq 0 $((PEERS - 1))); do
|
||||
PORT=$((BASE_PORT + i))
|
||||
ARGS=(--mode p2p --role peer --id "$i" --port "$PORT"
|
||||
--rounds "$ROUNDS" --trainers "$PEERS" --epochs "$EPOCHS")
|
||||
for j in $(seq 0 $((PEERS - 1))); do
|
||||
[ "$j" = "$i" ] && continue
|
||||
ARGS+=(--peer "localhost:$((BASE_PORT + j))")
|
||||
done
|
||||
echo "==> Starting peer-$i on port $PORT"
|
||||
"$BIN" "${ARGS[@]}" "${DATA_ARGS[@]}" > "$OUT/peer$i.log" 2>&1 &
|
||||
PIDS+=($!)
|
||||
done
|
||||
|
||||
echo "==> Federation running; waiting for all $PEERS peers to finish..."
|
||||
echo " (live log: tail -f $OUT/peer0.log)"
|
||||
FAIL=0
|
||||
for pid in "${PIDS[@]}"; do
|
||||
wait "$pid" || FAIL=1
|
||||
done
|
||||
PIDS=()
|
||||
if [ "$FAIL" -ne 0 ]; then
|
||||
echo "!! A peer exited abnormally; check $OUT/ logs"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo
|
||||
echo "==> Evaluation results (peer-0's view):"
|
||||
grep "Eval round=" "$OUT/peer0.log" || true
|
||||
echo
|
||||
|
||||
MODELS=(fl_model_peer-*.bin)
|
||||
if [ "${#MODELS[@]}" -ne "$PEERS" ]; then
|
||||
echo "!! Expected $PEERS model files, found ${#MODELS[@]}; check $OUT/ logs"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
FIRST="${MODELS[0]}"
|
||||
for m in "${MODELS[@]:1}"; do
|
||||
if ! cmp -s "$FIRST" "$m"; then
|
||||
echo "!! Divergence: $m differs from $FIRST"
|
||||
md5sum "${MODELS[@]}"
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
echo "==> Success: all $PEERS peers converged to a byte-identical model:"
|
||||
ls -la "${MODELS[@]}"
|
||||
echo " Full logs in $OUT/"
|
||||
Executable
+72
@@ -0,0 +1,72 @@
|
||||
#!/usr/bin/env bash
|
||||
#
|
||||
# Demo: federated learning in PROVIDER mode (star topology).
|
||||
# Starts one coordinator and N trainer nodes as separate processes, runs the
|
||||
# federation, then prints the evaluation log and the saved model file.
|
||||
#
|
||||
# Usage:
|
||||
# scripts/demo-provider.sh [ROUNDS] [TRAINERS] [EPOCHS] [/path/to/paysim.csv]
|
||||
#
|
||||
# Without a CSV path, nodes fall back to synthetic data (fast; the metrics are
|
||||
# meaningless, but the message flow and aggregation are fully exercised).
|
||||
set -euo pipefail
|
||||
|
||||
ROUNDS="${1:-3}"
|
||||
TRAINERS="${2:-3}"
|
||||
EPOCHS="${3:-2}"
|
||||
DATA="${4:-}"
|
||||
|
||||
BASE_PORT=9500
|
||||
OUT="demo-out/provider"
|
||||
|
||||
cd "$(dirname "$0")/.."
|
||||
echo "==> Building..."
|
||||
cabal build exe:fl-actors >/dev/null
|
||||
BIN="$(cabal list-bin fl-actors)"
|
||||
|
||||
mkdir -p "$OUT"
|
||||
rm -f "$OUT"/*.log fl_model_coordinator.bin
|
||||
|
||||
DATA_ARGS=()
|
||||
[ -n "$DATA" ] && DATA_ARGS=(--data "$DATA")
|
||||
|
||||
PIDS=()
|
||||
cleanup() { kill "${PIDS[@]}" 2>/dev/null || true; wait 2>/dev/null || true; }
|
||||
trap cleanup EXIT INT TERM
|
||||
|
||||
echo "==> Starting coordinator on port $BASE_PORT (rounds=$ROUNDS, trainers=$TRAINERS, epochs=$EPOCHS)"
|
||||
"$BIN" --mode provider --role coordinator \
|
||||
--port "$BASE_PORT" --rounds "$ROUNDS" --trainers "$TRAINERS" --epochs "$EPOCHS" \
|
||||
"${DATA_ARGS[@]}" > "$OUT/coordinator.log" 2>&1 &
|
||||
COORD=$!
|
||||
PIDS+=("$COORD")
|
||||
sleep 1
|
||||
|
||||
for i in $(seq 0 $((TRAINERS - 1))); do
|
||||
echo "==> Starting trainer-$i on port $((BASE_PORT + i + 1))"
|
||||
"$BIN" --mode provider --role trainer --id "$i" \
|
||||
--base-port "$BASE_PORT" --coord-host localhost --coord-port "$BASE_PORT" \
|
||||
--rounds "$ROUNDS" --trainers "$TRAINERS" --epochs "$EPOCHS" \
|
||||
"${DATA_ARGS[@]}" > "$OUT/trainer$i.log" 2>&1 &
|
||||
PIDS+=($!)
|
||||
done
|
||||
|
||||
echo "==> Federation running; waiting for the coordinator to finish..."
|
||||
echo " (live log: tail -f $OUT/coordinator.log)"
|
||||
if ! wait "$COORD"; then
|
||||
echo "!! Coordinator exited abnormally; last log lines:"
|
||||
tail -20 "$OUT/coordinator.log"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo
|
||||
echo "==> Evaluation results:"
|
||||
grep "Eval round=" "$OUT/coordinator.log" || true
|
||||
echo
|
||||
if [ -f fl_model_coordinator.bin ]; then
|
||||
echo "==> Success: model saved to fl_model_coordinator.bin ($(stat -c %s fl_model_coordinator.bin) bytes)"
|
||||
echo " Full logs in $OUT/"
|
||||
else
|
||||
echo "!! No model file was produced; check $OUT/ logs"
|
||||
exit 1
|
||||
fi
|
||||
@@ -5,6 +5,8 @@ module Control.Actor
|
||||
, module Control.Actor.Core
|
||||
, module Control.Actor.Supervision
|
||||
, module Control.Actor.Network
|
||||
-- Initialization
|
||||
, initRuntime
|
||||
-- Demo
|
||||
, pingActor
|
||||
, forwardActorWithCell
|
||||
|
||||
@@ -31,7 +31,7 @@ module Control.Actor.Core
|
||||
|
||||
import Control.Actor.Runtime
|
||||
import Control.Actor.Types
|
||||
import Control.Concurrent (forkIO, threadDelay)
|
||||
import Control.Concurrent (forkIO, killThread, threadDelay)
|
||||
import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
|
||||
import Control.Concurrent.STM
|
||||
( atomically
|
||||
@@ -52,8 +52,9 @@ import Control.Exception
|
||||
, throwIO
|
||||
, try
|
||||
)
|
||||
import Control.Monad (forM_, void)
|
||||
import Control.Monad (forM_, join, void)
|
||||
import Data.Maybe (fromMaybe)
|
||||
import System.Timeout (timeout)
|
||||
import Control.Monad.Reader
|
||||
( MonadIO (..)
|
||||
, MonadReader (..)
|
||||
@@ -122,6 +123,7 @@ spawnActorAs ::
|
||||
u ->
|
||||
RuntimeM (ActorRef m r)
|
||||
spawnActorAs uuid actorFn deathFn initState = do
|
||||
liftIO $ putStrLn $ "spawning actor with UUID " <> show uuid
|
||||
rt <- ask
|
||||
mailbox <- liftIO newTQueueIO
|
||||
deathQ <- liftIO newTQueueIO
|
||||
@@ -151,7 +153,13 @@ spawnActorAs uuid actorFn deathFn initState = do
|
||||
Stop -> return ()
|
||||
Continue u -> loop fn as {asEnv = u}
|
||||
|
||||
-- The thread blocks on `ready` until it is registered in rtActors:
|
||||
-- otherwise messages dispatched by UUID in the window before the insert
|
||||
-- would be dropped, and a death racing the insert would leave a stale
|
||||
-- entry behind (the death path's delete running before the insert).
|
||||
ready <- liftIO newEmptyMVar
|
||||
tid <- liftIO $ forkIO $ do
|
||||
takeMVar ready
|
||||
result <- try @SomeException (loop actorFn actorState)
|
||||
let reason = case result of
|
||||
Right () -> Normal
|
||||
@@ -167,6 +175,7 @@ spawnActorAs uuid actorFn deathFn initState = do
|
||||
|
||||
liftIO $ atomically $
|
||||
modifyTVar (rtActors rt) (Map.insert actorId (tid, SomeActorRef actorRef))
|
||||
liftIO $ putMVar ready ()
|
||||
return actorRef
|
||||
|
||||
spawnActor ::
|
||||
@@ -189,9 +198,16 @@ linkTo target = do
|
||||
as <- asks fst
|
||||
liftIO $ atomically $ modifyTVar (asLinks as) (target :)
|
||||
|
||||
-- | Forcibly terminate an actor's thread. Previously this only posted a
|
||||
-- 'DeathMessage' to the target's death queue, which a deathFn returning
|
||||
-- 'Continue' would silently ignore — and even on 'Stop' the links were told
|
||||
-- the exit was 'Normal'. Killing the thread makes links see 'Killed'.
|
||||
killActor :: ActorRef m r -> ActorM u ()
|
||||
killActor (LocalRef {arDeathQ, arState}) =
|
||||
liftIO $ atomically $ writeTQueue arDeathQ (DeathMessage (asId arState) Killed)
|
||||
killActor ref@(LocalRef {}) = do
|
||||
rt <- asks snd
|
||||
liftIO $ do
|
||||
actors <- readTVarIO (rtActors rt)
|
||||
forM_ (Map.lookup (actorRefId ref) actors) (killThread . fst)
|
||||
killActor (RemoteRef _) =
|
||||
liftIO $ putStrLn "killActor: remote kill not yet implemented"
|
||||
|
||||
@@ -206,13 +222,23 @@ cast' msg (RemoteRef (ActorId nodeId uuid)) = do
|
||||
maybeAddr <- lookupNode nodeId
|
||||
case maybeAddr of
|
||||
Nothing -> liftIO $ putStrLn $ "cast: no node in lookup table with id " <> show nodeId
|
||||
Just addr -> rtSendRemote rt addr (NMCast uuid (encode msg))
|
||||
Just addr -> do
|
||||
result <- liftIO $ try @SomeException $ withRuntime rt $ rtSendRemote rt addr (NMCast uuid (encode msg))
|
||||
case result of
|
||||
Left e -> liftIO $ putStrLn $ "cast: send to " <> show addr <> " failed: " <> show e
|
||||
Right _ -> return ()
|
||||
|
||||
-- | How long a call waits for a reply before giving up with Nothing.
|
||||
-- Without this, a call to a dead/missing actor (or over a dead connection)
|
||||
-- would block the caller forever.
|
||||
callTimeoutMicros :: Int
|
||||
callTimeoutMicros = 10_000_000
|
||||
|
||||
call' :: (Binary msg, Binary reply) => msg -> ActorRef msg reply -> RuntimeM (Maybe reply)
|
||||
call' msg (LocalRef {arMsgQ}) = liftIO $ do
|
||||
mv <- newEmptyMVar
|
||||
atomically $ writeTQueue arMsgQ (thisNodeId, Call msg mv)
|
||||
takeMVar mv
|
||||
join <$> timeout callTimeoutMicros (takeMVar mv)
|
||||
call' msg (RemoteRef (ActorId nodeId uuid)) = do
|
||||
rt <- ask
|
||||
corrId <- liftIO $ atomically $ do
|
||||
@@ -229,9 +255,9 @@ call' msg (RemoteRef (ActorId nodeId uuid)) = do
|
||||
return Nothing
|
||||
Just addr -> do
|
||||
rtSendRemote rt addr (NMCall uuid corrId (rtNodeId rt) (encode msg))
|
||||
raw <- liftIO $ takeMVar replyVar
|
||||
raw <- liftIO $ timeout callTimeoutMicros (takeMVar replyVar)
|
||||
liftIO $ atomically $ modifyTVar (rtPending rt) (Map.delete corrId)
|
||||
return $ Just (decode raw)
|
||||
return $ decode <$> raw
|
||||
|
||||
cast :: (Binary msg) => msg -> ActorRef msg reply -> ActorM u ()
|
||||
cast = (liftRuntime .) . cast'
|
||||
|
||||
+115
-73
@@ -15,9 +15,9 @@ module Control.Actor.Network
|
||||
|
||||
import Control.Actor.Core
|
||||
( ActorM, ActorResult (..), Handler, cast', liftRuntime, linkActorTo
|
||||
, spawnActor, state, stopOnDeath )
|
||||
, pass, spawnActor, state, stopOnDeath )
|
||||
import Control.Actor.Registry (deregisterNode)
|
||||
import Control.Actor.Runtime (Runtime (..), RuntimeM, addrToNodeId, withRuntime)
|
||||
import Control.Actor.Supervision (ChildSpec (..), RestartStrategy (..), childWithRef, supervise')
|
||||
import Control.Actor.Transport (ConnHandle (..), Transport (..), createTCPTransport)
|
||||
import Control.Actor.Types
|
||||
import Control.Concurrent (forkIO)
|
||||
@@ -26,23 +26,21 @@ import Control.Concurrent.STM
|
||||
( atomically
|
||||
, modifyTVar
|
||||
, newEmptyTMVar
|
||||
, newEmptyTMVarIO
|
||||
, putTMVar
|
||||
, readTMVar
|
||||
, readTVar
|
||||
, readTVarIO
|
||||
, tryPutTMVar
|
||||
, tryTakeTMVar
|
||||
, writeTQueue
|
||||
, writeTVar
|
||||
)
|
||||
import Control.Exception (SomeException, try)
|
||||
import Control.Monad (forM_, forever, unless, void)
|
||||
import Control.Monad.Reader (MonadIO (..), MonadReader (..))
|
||||
import Control.Exception (SomeException, throwIO, toException, try)
|
||||
import Control.Monad (forM_, forever, unless, void, when)
|
||||
import Control.Monad.Reader (MonadIO (..), MonadReader (..), asks)
|
||||
import Data.Binary (decode, decodeOrFail, encode)
|
||||
import Data.ByteString.Lazy (ByteString)
|
||||
import Data.Map qualified as Map
|
||||
import Data.Maybe (fromMaybe)
|
||||
import System.Timeout (timeout)
|
||||
|
||||
-- Connection actors
|
||||
|
||||
@@ -54,51 +52,63 @@ connActorFn nm = do
|
||||
|
||||
connDeathFn :: NodeAddr -> DeathMessage -> ActorM ConnHandle (SupervisorAction ConnHandle)
|
||||
connDeathFn peer _ = do
|
||||
ch <- state
|
||||
rt <- liftRuntime ask
|
||||
liftIO $ do
|
||||
atomically $ modifyTVar (rtConnections rt) (Map.delete peer)
|
||||
chClose ch
|
||||
ch <- state
|
||||
aid <- asks (asId . fst)
|
||||
rt <- liftRuntime ask
|
||||
-- Only clear the pool entry if it is ours: with simultaneous connects two
|
||||
-- conn actors can exist for the same peer, and the loser must not evict
|
||||
-- the active one.
|
||||
wasActive <- liftIO $ atomically $ do
|
||||
conns <- readTVar (rtConnections rt)
|
||||
case Map.lookup peer conns of
|
||||
Just ref | actorRefId ref == aid -> do
|
||||
modifyTVar (rtConnections rt) (Map.delete peer)
|
||||
return True
|
||||
_else -> return False
|
||||
liftIO $ chClose ch
|
||||
when wasActive $ do
|
||||
mNodeId <- liftRuntime $ addrToNodeId peer
|
||||
forM_ mNodeId $ \nodeId -> liftRuntime $ deregisterNode nodeId
|
||||
return Stop
|
||||
|
||||
routerActorFn :: NodeAddr -> Handler ByteString () ()
|
||||
routerActorFn senderAddr raw = do
|
||||
let nm = decode raw :: NetworkMessage
|
||||
valid <- liftRuntime $ validateAndDispatch senderAddr nm
|
||||
unless valid $ liftIO $ putStrLn "router: dropping invalid message"
|
||||
return $ ActorResult Nothing () Nothing
|
||||
routerActorFn senderAddr raw =
|
||||
case decodeOrFail raw of
|
||||
Left _ -> do
|
||||
liftIO $ putStrLn $ "router: dropping undecodable frame from " <> show senderAddr
|
||||
pass
|
||||
Right (_, _, nm) -> do
|
||||
valid <- liftRuntime $ validateAndDispatch senderAddr nm
|
||||
unless valid $ liftIO $
|
||||
putStrLn $ "router: dropping invalid message from " <> show senderAddr
|
||||
pass
|
||||
|
||||
-- Connection supervision tree
|
||||
|
||||
spawnConnTree :: NodeAddr -> ConnHandle -> RuntimeM (ActorRef NetworkMessage ())
|
||||
spawnConnTree peer ch = do
|
||||
rt <- ask
|
||||
routerCell <- liftIO newEmptyTMVarIO
|
||||
connCell <- liftIO newEmptyTMVarIO
|
||||
supervise' OneForAll
|
||||
[ childWithRef (routerActorFn peer) stopOnDeath () routerCell
|
||||
, ChildSpec
|
||||
{ csRun = \target -> do
|
||||
ref <- spawnActor connActorFn (connDeathFn peer) ch
|
||||
linkActorTo target ref
|
||||
return ref
|
||||
, csOnSpawn = \ref -> case ref of
|
||||
LocalRef { arDeathQ, arState } -> do
|
||||
atomically $ do
|
||||
void $ tryTakeTMVar connCell
|
||||
putTMVar connCell ref
|
||||
routerRef <- atomically $ readTMVar routerCell
|
||||
void $ forkIO $ do
|
||||
result <- try @SomeException $ forever $
|
||||
chRecv ch >>= \raw -> withRuntime rt $ cast' raw routerRef
|
||||
case result of
|
||||
Left _ -> atomically $
|
||||
writeTQueue arDeathQ (DeathMessage (asId arState) Killed)
|
||||
Right () -> return ()
|
||||
RemoteRef _ -> return ()
|
||||
}
|
||||
]
|
||||
liftIO $ atomically $ readTMVar connCell
|
||||
rt <- ask
|
||||
routerRef <- spawnActor (routerActorFn peer) stopOnDeath ()
|
||||
connRef <- spawnActor connActorFn (connDeathFn peer) ch
|
||||
-- When conn dies (peer disconnected), notify router so it shuts down too
|
||||
case routerRef of
|
||||
LocalRef {arDeathQ} -> linkActorTo (LocalTarget arDeathQ) connRef
|
||||
RemoteRef _ -> return ()
|
||||
-- And vice versa: if the router dies, close the connection down cleanly
|
||||
case connRef of
|
||||
LocalRef {arDeathQ} -> linkActorTo (LocalTarget arDeathQ) routerRef
|
||||
RemoteRef _ -> return ()
|
||||
liftIO $ case connRef of
|
||||
LocalRef {arDeathQ, arState} ->
|
||||
void $ forkIO $ do
|
||||
result <- try @SomeException $ forever $
|
||||
chRecv ch >>= \raw -> withRuntime rt $ cast' raw routerRef
|
||||
case result of
|
||||
Left _ -> atomically $
|
||||
writeTQueue arDeathQ (DeathMessage (asId arState) Killed)
|
||||
Right () -> return ()
|
||||
RemoteRef _ -> return ()
|
||||
return connRef
|
||||
|
||||
-- Incoming connection handler
|
||||
|
||||
@@ -110,13 +120,26 @@ handleNewConn ch = do
|
||||
NMHandshake peerAddr -> do
|
||||
ref <- spawnConnTree peerAddr ch
|
||||
liftIO $ atomically $ do
|
||||
modifyTVar (rtConnections rt) (Map.insert peerAddr ref)
|
||||
conns <- readTVar (rtConnections rt)
|
||||
-- On a simultaneous connect two sockets (and two conn trees) exist
|
||||
-- for the same peer. Only the pool entry is deduplicated; the extra
|
||||
-- tree must stay alive because it is the read path for whatever the
|
||||
-- peer sends on its socket. It cleans itself up when that socket
|
||||
-- closes. Local waiters get the pool winner, never the extra ref.
|
||||
let winner = Map.findWithDefault ref peerAddr conns
|
||||
unless (Map.member peerAddr conns) $
|
||||
modifyTVar (rtConnections rt) (Map.insert peerAddr ref)
|
||||
table <- readTVar (rtNodeTable rt)
|
||||
unless (any (== peerAddr) (Map.elems table)) $ do
|
||||
nid <- readTVar (rtNextNodeId rt)
|
||||
writeTVar (rtNextNodeId rt) (nid + 1)
|
||||
modifyTVar (rtNodeTable rt) (Map.insert nid peerAddr)
|
||||
promises <- readTVar (rtConnPromises rt)
|
||||
case Map.lookup peerAddr promises of
|
||||
Nothing -> return ()
|
||||
Just p -> do
|
||||
modifyTVar (rtConnPromises rt) (Map.delete peerAddr)
|
||||
void $ tryPutTMVar p ref
|
||||
void $ tryPutTMVar p (Right winner)
|
||||
_else -> liftIO $ chClose ch
|
||||
|
||||
-- Connection pool
|
||||
@@ -138,18 +161,29 @@ getOrCreateConn peer = do
|
||||
return (Right (Right p))
|
||||
case action of
|
||||
Left ref -> return ref
|
||||
Right (Left promise) -> liftIO $ atomically $ readTMVar promise
|
||||
Right (Left promise) -> do
|
||||
-- Someone else is connecting; wait for their outcome so a failed
|
||||
-- connect propagates instead of blocking us forever.
|
||||
outcome <- liftIO $ atomically $ readTMVar promise
|
||||
either (liftIO . throwIO) return outcome
|
||||
Right (Right promise) -> do
|
||||
ch <- liftIO $ tConnect (rtTransport rt) peer
|
||||
liftIO $ chSend ch (encode (NMHandshake (rtNodeId rt)))
|
||||
ref <- spawnConnTree peer ch
|
||||
liftIO $ atomically $ do
|
||||
conns <- readTVar (rtConnections rt)
|
||||
unless (Map.member peer conns) $
|
||||
modifyTVar (rtConnections rt) (Map.insert peer ref)
|
||||
modifyTVar (rtConnPromises rt) (Map.delete peer)
|
||||
void $ tryPutTMVar promise ref
|
||||
return ref
|
||||
result <- liftIO $ try @SomeException $ tConnect (rtTransport rt) peer
|
||||
case result of
|
||||
Left e -> do
|
||||
liftIO $ atomically $ do
|
||||
modifyTVar (rtConnPromises rt) (Map.delete peer)
|
||||
void $ tryPutTMVar promise (Left e)
|
||||
liftIO $ throwIO e
|
||||
Right ch -> do
|
||||
liftIO $ chSend ch (encode (NMHandshake (rtNodeId rt)))
|
||||
ref <- spawnConnTree peer ch
|
||||
liftIO $ atomically $ do
|
||||
conns <- readTVar (rtConnections rt)
|
||||
unless (Map.member peer conns) $
|
||||
modifyTVar (rtConnections rt) (Map.insert peer ref)
|
||||
modifyTVar (rtConnPromises rt) (Map.delete peer)
|
||||
void $ tryPutTMVar promise (Right ref)
|
||||
return ref
|
||||
|
||||
-- Message dispatch
|
||||
|
||||
@@ -162,7 +196,7 @@ routeRemoteDeath senderAddr (ActorId _ uuid) reason = do
|
||||
exitReason = case reason of
|
||||
RNormal -> Normal
|
||||
RKilled -> Killed
|
||||
RException s -> Exception (error s)
|
||||
RException s -> Exception (toException (userError s))
|
||||
dm = DeathMessage deadId exitReason
|
||||
actors <- readTVarIO (rtActors rt)
|
||||
forM_ (Map.elems actors) $ \(_, SomeActorRef ref) ->
|
||||
@@ -195,7 +229,8 @@ validateAndDispatch senderAddr nm = case nm of
|
||||
liftIO $ do
|
||||
actors <- readTVarIO (rtActors rt)
|
||||
case findByUUID uuid actors of
|
||||
Nothing ->
|
||||
Nothing -> do
|
||||
putStrLn $ "didn't find actor: " <> show uuid
|
||||
return False
|
||||
Just (SomeActorRef (LocalRef {arMsgQ})) ->
|
||||
case decodeOrFail payload of
|
||||
@@ -221,12 +256,16 @@ validateAndDispatch senderAddr nm = case nm of
|
||||
mv <- newEmptyMVar
|
||||
atomically $ writeTQueue arMsgQ (rid, Call msg mv)
|
||||
void $ forkIO $ do
|
||||
reply <- takeMVar mv
|
||||
-- Bounded wait: if the target actor dies before replying,
|
||||
-- this thread must not linger forever.
|
||||
reply <- timeout 30_000_000 (takeMVar mv)
|
||||
case reply of
|
||||
Nothing -> return ()
|
||||
Just rv -> withRuntime rt $ do
|
||||
connRef <- getOrCreateConn returnAddr
|
||||
cast' (NMReply corrId (encode rv)) connRef
|
||||
Nothing -> return ()
|
||||
Just Nothing -> return ()
|
||||
Just (Just rv) -> void $ try @SomeException $
|
||||
withRuntime rt $ do
|
||||
connRef <- getOrCreateConn returnAddr
|
||||
cast' (NMReply corrId (encode rv)) connRef
|
||||
return True
|
||||
Just (SomeActorRef (RemoteRef _)) ->
|
||||
return False
|
||||
@@ -252,13 +291,16 @@ connect :: Maybe NodeId -> NodeAddr -> RuntimeM NodeId
|
||||
connect suggestedId peer = do
|
||||
rt <- ask
|
||||
nodeId <- liftIO $ atomically $ do
|
||||
table <- readTVar (rtNodeTable rt)
|
||||
nid <- readTVar (rtNextNodeId rt)
|
||||
let assigned = case suggestedId of
|
||||
Just n | n /= 0 && not (Map.member n table) -> n
|
||||
_else -> nid
|
||||
writeTVar (rtNextNodeId rt) (max (assigned + 1) nid)
|
||||
modifyTVar (rtNodeTable rt) (Map.insert assigned peer)
|
||||
return assigned
|
||||
table <- readTVar (rtNodeTable rt)
|
||||
case Map.foldrWithKey (\k v acc -> if v == peer then Just k else acc) Nothing table of
|
||||
Just existing -> return existing
|
||||
Nothing -> do
|
||||
nid <- readTVar (rtNextNodeId rt)
|
||||
let assigned = case suggestedId of
|
||||
Just n | n /= 0 && not (Map.member n table) -> n
|
||||
_else -> nid
|
||||
writeTVar (rtNextNodeId rt) (max (assigned + 1) nid)
|
||||
modifyTVar (rtNodeTable rt) (Map.insert assigned peer)
|
||||
return assigned
|
||||
void $ getOrCreateConn peer
|
||||
return nodeId
|
||||
|
||||
@@ -5,13 +5,14 @@ module Control.Actor.Registry
|
||||
registryUUID,
|
||||
registry,
|
||||
registerDeath,
|
||||
deregisterNode,
|
||||
createRegistry,
|
||||
)
|
||||
where
|
||||
|
||||
import Control.Actor.Core (ActorM, ActorResult (..), Handler, call, call', cast', liftRuntime, linkActorTo, pass, spawnActorAs, state, lastMessageFrom, cast, passWith)
|
||||
import Control.Actor.Runtime (Runtime (..), RuntimeM, getActorByUUID, withRuntime)
|
||||
import Control.Actor.Types (ActorId (..), ActorRef (..), DeathMessage (..), DeathTarget (LocalTarget), RegistryMsg (..), SomeActorRef (..), SupervisorAction (Continue), actorRefId, thisNodeId, findByUUID)
|
||||
import Control.Actor.Types (ActorId (..), ActorRef (..), DeathMessage (..), DeathTarget (LocalTarget), NodeId, RegistryMsg (..), SomeActorRef (..), SupervisorAction (Continue), actorRefId, thisNodeId, findByUUID)
|
||||
import Control.Concurrent.STM (readTVarIO)
|
||||
import Control.Monad (forM_, when)
|
||||
import Control.Monad.IO.Class (MonadIO (liftIO))
|
||||
@@ -41,6 +42,9 @@ registryHandlerFn (RMLookup name) = do
|
||||
registryHandlerFn (RMDeregister name) = do
|
||||
u <- state
|
||||
return $ ActorResult Nothing (Map.delete name u) Nothing
|
||||
registryHandlerFn (RMDeregisterNode nodeId) = do
|
||||
u <- state
|
||||
passWith (Map.filter (\(ActorId nid _) -> nid /= nodeId) u)
|
||||
registryHandlerFn msg@(RMDeath deadUUID) = do
|
||||
u <- state
|
||||
from <- lastMessageFrom
|
||||
@@ -68,13 +72,21 @@ registryDeathFn (DeathMessage (ActorId _ deadUUID) _) = do
|
||||
|
||||
registry :: RuntimeM (Maybe (ActorRef RegistryMsg (Maybe ActorId)))
|
||||
registry = do
|
||||
unsafeCoerce $ getActorByUUID registryUUID
|
||||
mRef <- getActorByUUID registryUUID
|
||||
return $ case mRef of
|
||||
Nothing -> Nothing
|
||||
Just (SomeActorRef regRef) -> Just (unsafeCoerce regRef)
|
||||
|
||||
registerDeath :: ActorId -> RuntimeM ()
|
||||
registerDeath (ActorId _ uuid) = do
|
||||
r <- registry
|
||||
mapM_ (cast' (RMDeath uuid)) r
|
||||
|
||||
deregisterNode :: NodeId -> RuntimeM ()
|
||||
deregisterNode nodeId = do
|
||||
r <- registry
|
||||
mapM_ (cast' (RMDeregisterNode nodeId)) r
|
||||
|
||||
registerActor :: String -> ActorRef msg r -> RuntimeM ()
|
||||
registerActor name ref = do
|
||||
rt <- ask
|
||||
@@ -122,8 +134,6 @@ lookupRemoteActor name = do
|
||||
case result of
|
||||
Just (Just (ActorId 0 uuid)) ->
|
||||
return $ Just (RemoteRef (ActorId peerNodeId uuid))
|
||||
Just (Just aid) ->
|
||||
return $ Just (RemoteRef aid)
|
||||
_else -> go rest
|
||||
|
||||
createRegistry :: RuntimeM (ActorRef RegistryMsg (Maybe ActorId))
|
||||
|
||||
@@ -14,6 +14,7 @@ module Control.Actor.Runtime
|
||||
import Control.Actor.Transport (Transport)
|
||||
import Control.Actor.Types
|
||||
import Control.Concurrent (ThreadId)
|
||||
import Control.Exception (SomeException)
|
||||
import Control.Concurrent.MVar (MVar)
|
||||
import Control.Concurrent.STM
|
||||
( TMVar, TVar, atomically, newTVarIO, readTVar, readTVarIO )
|
||||
@@ -33,7 +34,7 @@ data Runtime = Runtime
|
||||
, rtNodeTable :: TVar (Map.Map NodeId NodeAddr)
|
||||
, rtTransport :: Transport
|
||||
, rtConnections :: TVar (Map.Map NodeAddr (ActorRef NetworkMessage ()))
|
||||
, rtConnPromises :: TVar (Map.Map NodeAddr (TMVar (ActorRef NetworkMessage ())))
|
||||
, rtConnPromises :: TVar (Map.Map NodeAddr (TMVar (Either SomeException (ActorRef NetworkMessage ()))))
|
||||
, rtSendRemote :: NodeAddr -> NetworkMessage -> RuntimeM ()
|
||||
}
|
||||
|
||||
|
||||
@@ -23,10 +23,8 @@ import Data.Binary (Binary)
|
||||
import Control.Concurrent (ThreadId, forkIO, killThread)
|
||||
import Control.Concurrent.STM
|
||||
( TMVar
|
||||
, TQueue
|
||||
, TVar
|
||||
, atomically
|
||||
, flushTQueue
|
||||
, modifyTVar
|
||||
, newTQueueIO
|
||||
, newTVarIO
|
||||
@@ -36,13 +34,16 @@ import Control.Concurrent.STM
|
||||
, tryTakeTMVar
|
||||
, writeTVar
|
||||
)
|
||||
import Control.Monad (forever, void)
|
||||
import Control.Monad (void)
|
||||
import GHC.Clock (getMonotonicTime)
|
||||
import Control.Monad.Reader (MonadIO (..), MonadReader (..))
|
||||
import Data.Map qualified as Map
|
||||
|
||||
data ChildSpec = forall m r. (Binary m, Binary r) => ChildSpec
|
||||
{ csRun :: DeathTarget -> RuntimeM (ActorRef m r)
|
||||
, csOnSpawn :: ActorRef m r -> IO ()
|
||||
-- | Runs after every (re)spawn — including supervisor restarts — so it can
|
||||
-- refresh shared ref cells and re-register names.
|
||||
, csOnSpawn :: ActorRef m r -> RuntimeM ()
|
||||
}
|
||||
|
||||
child ::
|
||||
@@ -73,7 +74,7 @@ childWithRef msgFn deathFn initState cell =
|
||||
ref <- spawnActor msgFn deathFn initState
|
||||
linkActorTo target ref
|
||||
return ref
|
||||
, csOnSpawn = \ref -> atomically $ do
|
||||
, csOnSpawn = \ref -> liftIO $ atomically $ do
|
||||
void $ tryTakeTMVar cell
|
||||
putTMVar cell ref
|
||||
}
|
||||
@@ -91,7 +92,7 @@ spawnSlot :: DeathTarget -> ChildSpec -> RuntimeM ChildSlot
|
||||
spawnSlot target spec@ChildSpec {csRun, csOnSpawn} = do
|
||||
rt <- ask
|
||||
ref <- csRun target
|
||||
_ <- liftIO $ csOnSpawn ref
|
||||
csOnSpawn ref
|
||||
let aid = someActorId (SomeActorRef ref)
|
||||
tid <- liftIO $ do
|
||||
actors <- readTVarIO (rtActors rt)
|
||||
@@ -103,6 +104,15 @@ spawnSlot target spec@ChildSpec {csRun, csOnSpawn} = do
|
||||
killSlot :: ChildSlot -> IO ()
|
||||
killSlot ChildSlot {slotTid} = killThread slotTid
|
||||
|
||||
-- | Max child restarts within 'restartWindow' seconds before the supervisor
|
||||
-- gives up and stops all children — otherwise a child that crashes right
|
||||
-- after spawning would respawn in a hot loop forever.
|
||||
restartLimit :: Int
|
||||
restartLimit = 5
|
||||
|
||||
restartWindow :: Double
|
||||
restartWindow = 10
|
||||
|
||||
supervise' :: RestartStrategy -> [ChildSpec] -> RuntimeM ()
|
||||
supervise' strategy specs = do
|
||||
rt <- ask
|
||||
@@ -110,13 +120,35 @@ supervise' strategy specs = do
|
||||
let target = LocalTarget supDeathQ
|
||||
slots <- mapM (spawnSlot target) specs
|
||||
slotsVar <- liftIO $ newTVarIO slots
|
||||
liftIO $ void $ forkIO $ forever $ do
|
||||
DeathMessage deadId _ <- atomically $ readTQueue supDeathQ
|
||||
slots' <- readTVarIO slotsVar
|
||||
withRuntime rt $ case strategy of
|
||||
OneForOne -> doOneForOne target slotsVar slots' deadId
|
||||
OneForAll -> doOneForAll target slotsVar supDeathQ slots'
|
||||
RestForOne -> doRestForOne target slotsVar supDeathQ slots' deadId
|
||||
let loop restarts = do
|
||||
DeathMessage deadId reason <- atomically $ readTQueue supDeathQ
|
||||
slots' <- readTVarIO slotsVar
|
||||
if all ((/= deadId) . slotId) slots'
|
||||
-- Stale death: killSlot is asynchronous, so children replaced by a
|
||||
-- OneForAll/RestForOne restart report their deaths after the fact.
|
||||
-- Acting on one would kill the fresh children and cascade forever.
|
||||
then loop restarts
|
||||
else case reason of
|
||||
-- A normal exit is completion, not a crash: drop the slot.
|
||||
Normal -> do
|
||||
atomically $ modifyTVar slotsVar (filter ((/= deadId) . slotId))
|
||||
loop restarts
|
||||
_crash -> do
|
||||
now <- getMonotonicTime
|
||||
let recent = filter (> now - restartWindow) restarts
|
||||
if length recent >= restartLimit
|
||||
then do
|
||||
putStrLn $ "supervisor: " <> show restartLimit
|
||||
<> " restarts within " <> show restartWindow
|
||||
<> "s, giving up"
|
||||
mapM_ killSlot slots'
|
||||
else do
|
||||
withRuntime rt $ case strategy of
|
||||
OneForOne -> doOneForOne target slotsVar slots' deadId
|
||||
OneForAll -> doOneForAll target slotsVar slots'
|
||||
RestForOne -> doRestForOne target slotsVar slots' deadId
|
||||
loop (now : recent)
|
||||
liftIO $ void $ forkIO $ loop []
|
||||
|
||||
supervise :: RestartStrategy -> [ChildSpec] -> ActorM u ()
|
||||
supervise = (liftRuntime .) . supervise'
|
||||
@@ -135,30 +167,24 @@ doOneForOne target slotsVar slots deadId =
|
||||
doOneForAll ::
|
||||
DeathTarget ->
|
||||
TVar [ChildSlot] ->
|
||||
TQueue DeathMessage ->
|
||||
[ChildSlot] ->
|
||||
RuntimeM ()
|
||||
doOneForAll target slotsVar supDeathQ slots = do
|
||||
liftIO $ do
|
||||
mapM_ killSlot slots
|
||||
atomically $ void $ flushTQueue supDeathQ
|
||||
doOneForAll target slotsVar slots = do
|
||||
liftIO $ mapM_ killSlot slots
|
||||
newSlots <- mapM (spawnSlot target . slotSpec) slots
|
||||
liftIO $ atomically $ writeTVar slotsVar newSlots
|
||||
|
||||
doRestForOne ::
|
||||
DeathTarget ->
|
||||
TVar [ChildSlot] ->
|
||||
TQueue DeathMessage ->
|
||||
[ChildSlot] ->
|
||||
ActorId ->
|
||||
RuntimeM ()
|
||||
doRestForOne target slotsVar supDeathQ slots deadId = do
|
||||
doRestForOne target slotsVar slots deadId = do
|
||||
let (before, fromDead) = break (\s -> slotId s == deadId) slots
|
||||
case fromDead of
|
||||
[] -> return ()
|
||||
_nonempty -> do
|
||||
liftIO $ do
|
||||
mapM_ killSlot (drop 1 fromDead)
|
||||
atomically $ void $ flushTQueue supDeathQ
|
||||
liftIO $ mapM_ killSlot (drop 1 fromDead)
|
||||
newSlots <- mapM (spawnSlot target . slotSpec) fromDead
|
||||
liftIO $ atomically $ writeTVar slotsVar (before ++ newSlots)
|
||||
|
||||
@@ -85,7 +85,7 @@ data NetworkMessage
|
||||
| NMCall UUID CorrelationId NodeAddr ByteString
|
||||
| NMReply CorrelationId ByteString
|
||||
| NMDeath ActorId RemoteExitReason
|
||||
deriving (Generic)
|
||||
deriving (Generic, Show)
|
||||
|
||||
instance Binary NetworkMessage
|
||||
|
||||
@@ -125,10 +125,11 @@ data SupervisorAction u
|
||||
| Continue u
|
||||
|
||||
data RegistryMsg
|
||||
= RMRegister String UUID
|
||||
| RMLookup String
|
||||
| RMDeregister String
|
||||
| RMDeath UUID
|
||||
= RMRegister String UUID
|
||||
| RMLookup String
|
||||
| RMDeregister String
|
||||
| RMDeath UUID
|
||||
| RMDeregisterNode NodeId
|
||||
deriving (Generic)
|
||||
|
||||
instance Binary RegistryMsg
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
module FL.Actors.Aggregator
|
||||
( AggregatorState (..)
|
||||
, aggregatorHandler
|
||||
) where
|
||||
|
||||
import Control.Actor.Core
|
||||
import Control.Actor.Types
|
||||
import Control.Concurrent.STM (TMVar, atomically, readTMVar)
|
||||
import Control.Monad.IO.Class (liftIO)
|
||||
import Data.List (sortOn)
|
||||
import FL.Actors.Logger (logFn)
|
||||
import FL.NN (fedAvg)
|
||||
import FL.Types
|
||||
|
||||
data AggregatorState = AggregatorState
|
||||
{ asExpected :: Int
|
||||
, asUpdates :: [ModelUpdate]
|
||||
, asCoordCell :: TMVar (ActorRef CoordinatorMsg ()) -- filled after coordinator spawns
|
||||
, asLogRef :: ActorRef LogMsg ()
|
||||
, asCurrentRound :: Int
|
||||
}
|
||||
|
||||
aggregatorHandler :: Handler AggregatorMsg AggregatorState ()
|
||||
aggregatorHandler (SubmitUpdate upd) = do
|
||||
st <- state
|
||||
if muRound upd /= asCurrentRound st
|
||||
then pass
|
||||
else do
|
||||
let updates' = upd : filter (\u -> muTrainerId u /= muTrainerId upd) (asUpdates st)
|
||||
if length updates' >= asExpected st
|
||||
then do
|
||||
let pairs = [(muSampleSize u, wData (muWeights u)) | u <- sortOn muTrainerId updates']
|
||||
avgW = Weights (fedAvg pairs)
|
||||
nextRound = asCurrentRound st + 1
|
||||
newModel = GlobalModel nextRound avgW
|
||||
logFn (asLogRef st) "Aggregator" INFO $
|
||||
"Round " <> show (asCurrentRound st)
|
||||
<> ": aggregated " <> show (length updates') <> " updates"
|
||||
coordRef <- liftIO $ atomically $ readTMVar (asCoordCell st)
|
||||
cast (AggregationComplete newModel) coordRef
|
||||
passWith st { asUpdates = [], asCurrentRound = nextRound }
|
||||
else
|
||||
passWith st { asUpdates = updates' }
|
||||
@@ -0,0 +1,103 @@
|
||||
module FL.Actors.Coordinator
|
||||
( CoordinatorState (..)
|
||||
, coordinatorHandler
|
||||
, coordinatorDeathHandler
|
||||
) where
|
||||
|
||||
import Control.Actor.Core
|
||||
import Control.Actor.Types
|
||||
import Data.Binary (encodeFile)
|
||||
import Control.Concurrent.MVar (MVar, putMVar)
|
||||
import Control.Monad (forM_, when)
|
||||
import Control.Monad.IO.Class (liftIO)
|
||||
import Data.Map.Strict (Map)
|
||||
import Data.Map.Strict qualified as Map
|
||||
import Data.Set qualified as Set
|
||||
import Data.UUID.V4 (nextRandom)
|
||||
import FL.Actors.Logger (logFn)
|
||||
import FL.CRDT (ORSet, orAdd, orMembers)
|
||||
import FL.Types
|
||||
|
||||
data CoordinatorState = CoordinatorState
|
||||
{ csTrainers :: ORSet String
|
||||
, csTrainerRefs :: Map String (ActorRef TrainerMsg ())
|
||||
, csExpectedTrainers :: Int
|
||||
, csCurrentRound :: Int
|
||||
, csTotalRounds :: Int
|
||||
, csGlobalModel :: GlobalModel
|
||||
, csAggRef :: ActorRef AggregatorMsg ()
|
||||
, csEvalRef :: ActorRef EvaluatorMsg ()
|
||||
, csLogRef :: ActorRef LogMsg ()
|
||||
, csDoneVar :: MVar ()
|
||||
}
|
||||
|
||||
coordinatorHandler :: Handler CoordinatorMsg CoordinatorState ()
|
||||
coordinatorHandler (RegisterTrainer tid aid) = do
|
||||
st <- state
|
||||
-- The ActorId in the message carries the *sender's* node numbering (its own
|
||||
-- node id is always 0 to itself); rebase it onto our node id for the sender,
|
||||
-- otherwise every cast to the trainer is dropped ("no node with id 0").
|
||||
from <- lastMessageFrom
|
||||
uuid <- liftIO nextRandom
|
||||
let ActorId _ trainerUUID = aid
|
||||
ref = RemoteRef (ActorId from trainerUUID) :: ActorRef TrainerMsg ()
|
||||
trainers' = orAdd uuid tid (csTrainers st)
|
||||
refs' = Map.insert tid ref (csTrainerRefs st)
|
||||
st' = st { csTrainers = trainers', csTrainerRefs = refs' }
|
||||
numRegistered = Set.size (orMembers trainers')
|
||||
logFn (csLogRef st) "Coordinator" INFO $
|
||||
"Trainer registered: " <> tid
|
||||
<> " (" <> show numRegistered <> "/" <> show (csExpectedTrainers st) <> ")"
|
||||
-- Auto-start once all expected trainers have registered (per the OR-Set)
|
||||
if numRegistered >= csExpectedTrainers st
|
||||
then startRound st'
|
||||
else passWith st'
|
||||
|
||||
coordinatorHandler StartFederation = do
|
||||
st <- state
|
||||
startRound st
|
||||
|
||||
coordinatorHandler (AggregationComplete newModel) = do
|
||||
st <- state
|
||||
logFn (csLogRef st) "Coordinator" INFO $
|
||||
"Round " <> show (gmRound newModel - 1) <> " aggregation complete"
|
||||
cast (EvaluateModel newModel) (csEvalRef st)
|
||||
let st' = st { csGlobalModel = newModel, csCurrentRound = gmRound newModel }
|
||||
if csCurrentRound st' >= csTotalRounds st'
|
||||
then do
|
||||
liftIO $ encodeFile "fl_model_coordinator.bin" (gmWeights newModel)
|
||||
logFn (csLogRef st') "Coordinator" INFO
|
||||
"Model saved to fl_model_coordinator.bin"
|
||||
passWith st'
|
||||
else
|
||||
startRound st'
|
||||
|
||||
coordinatorHandler (EvaluationDone metrics) = do
|
||||
st <- state
|
||||
logFn (csLogRef st) "Coordinator" INFO $
|
||||
"Eval round=" <> show (mRound metrics)
|
||||
<> " acc=" <> show4 (mAccuracy metrics)
|
||||
<> " prec=" <> show4 (mPrecision metrics)
|
||||
<> " rec=" <> show4 (mRecall metrics)
|
||||
<> " f1=" <> show4 (mF1 metrics)
|
||||
<> " auc=" <> show4 (mAUCROC metrics)
|
||||
liftIO $ when (mRound metrics >= csTotalRounds st) $
|
||||
putMVar (csDoneVar st) ()
|
||||
pass
|
||||
where
|
||||
show4 x = show (fromIntegral (round (x * 10000) :: Int) / 10000.0 :: Double)
|
||||
|
||||
startRound :: CoordinatorState -> Actor CoordinatorMsg CoordinatorState ()
|
||||
startRound st = do
|
||||
let model = csGlobalModel st
|
||||
logFn (csLogRef st) "Coordinator" INFO $
|
||||
"Starting round " <> show (gmRound model + 1)
|
||||
<> " with " <> show (Map.size (csTrainerRefs st)) <> " trainers"
|
||||
forM_ (Map.elems (csTrainerRefs st)) $ \ref ->
|
||||
cast (StartTraining model) ref
|
||||
passWith st
|
||||
|
||||
coordinatorDeathHandler :: DeathMessage -> ActorM CoordinatorState (SupervisorAction CoordinatorState)
|
||||
coordinatorDeathHandler _ = do
|
||||
st <- state
|
||||
return (Continue st)
|
||||
@@ -0,0 +1,36 @@
|
||||
module FL.Actors.Evaluator
|
||||
( EvaluatorState (..)
|
||||
, evaluatorHandler
|
||||
) where
|
||||
|
||||
import Control.Actor.Core
|
||||
import Control.Actor.Types
|
||||
import Control.Concurrent.STM (TMVar, atomically, readTMVar)
|
||||
import Control.Monad.IO.Class (liftIO)
|
||||
import FL.Actors.Logger (logFn)
|
||||
import FL.NN (LocalData, evaluate, unflattenWeights)
|
||||
import FL.Types
|
||||
|
||||
data EvaluatorState = EvaluatorState
|
||||
{ esTestData :: LocalData
|
||||
, esCoordCell :: TMVar (ActorRef CoordinatorMsg ()) -- filled after coordinator spawns
|
||||
, esLogRef :: ActorRef LogMsg ()
|
||||
}
|
||||
|
||||
evaluatorHandler :: Handler EvaluatorMsg EvaluatorState ()
|
||||
evaluatorHandler (EvaluateModel gm) = do
|
||||
st <- state
|
||||
let ws = unflattenWeights (wData (gmWeights gm))
|
||||
metrics = evaluate ws (esTestData st) (gmRound gm)
|
||||
logFn (esLogRef st) "Evaluator" INFO $
|
||||
"Round " <> show (gmRound gm)
|
||||
<> " acc=" <> show4 (mAccuracy metrics)
|
||||
<> " prec=" <> show4 (mPrecision metrics)
|
||||
<> " rec=" <> show4 (mRecall metrics)
|
||||
<> " f1=" <> show4 (mF1 metrics)
|
||||
<> " auc=" <> show4 (mAUCROC metrics)
|
||||
coordRef <- liftIO $ atomically $ readTMVar (esCoordCell st)
|
||||
cast (EvaluationDone metrics) coordRef
|
||||
pass
|
||||
where
|
||||
show4 x = show (fromIntegral (round (x * 10000) :: Int) / 10000.0 :: Double)
|
||||
@@ -0,0 +1,37 @@
|
||||
module FL.Actors.Logger
|
||||
( loggerHandler
|
||||
, logFn
|
||||
, logFn'
|
||||
) where
|
||||
|
||||
import Control.Actor.Core (ActorM, Handler, cast, cast', pass)
|
||||
import Control.Actor.Runtime (RuntimeM)
|
||||
import Control.Actor.Types
|
||||
import Control.Monad.IO.Class (liftIO)
|
||||
import Data.Time.Clock (getCurrentTime)
|
||||
import FL.Types
|
||||
|
||||
type LogState = ()
|
||||
|
||||
loggerHandler :: Handler LogMsg LogState ()
|
||||
loggerHandler (Log LogEntry{leSource, leLevel, leMsg}) = do
|
||||
t <- liftIO getCurrentTime
|
||||
liftIO $ putStrLn $
|
||||
"[" <> show t <> "] [" <> showLevel leLevel <> "] "
|
||||
<> leSource <> ": " <> leMsg
|
||||
pass
|
||||
where
|
||||
showLevel DEBUG = "DEBUG"
|
||||
showLevel INFO = "INFO "
|
||||
showLevel WARN = "WARN "
|
||||
showLevel ERROR = "ERROR"
|
||||
|
||||
-- | Log via the logger actor. Goes through the normal cast path so it works
|
||||
-- for both local and remote logger refs (a remote ref used to be a silent
|
||||
-- no-op, losing all trainer-node logs in provider mode).
|
||||
logFn :: ActorRef LogMsg () -> String -> LogLevel -> String -> ActorM u ()
|
||||
logFn ref src lvl msg = cast (Log (LogEntry src lvl msg)) ref
|
||||
|
||||
-- | Same as 'logFn' but usable outside an actor (node startup code).
|
||||
logFn' :: ActorRef LogMsg () -> String -> LogLevel -> String -> RuntimeM ()
|
||||
logFn' ref src lvl msg = cast' (Log (LogEntry src lvl msg)) ref
|
||||
@@ -0,0 +1,166 @@
|
||||
{-# LANGUAGE StrictData #-}
|
||||
|
||||
module FL.Actors.PeerCoordinator
|
||||
( PeerCoordinatorState (..)
|
||||
, peerCoordinatorHandler
|
||||
, peerCoordinatorDeathHandler
|
||||
) where
|
||||
|
||||
import Control.Actor.Core
|
||||
import Control.Actor.Types
|
||||
import Data.Binary (encodeFile)
|
||||
import Unsafe.Coerce (unsafeCoerce)
|
||||
import Control.Concurrent.MVar (MVar, putMVar)
|
||||
import Control.Monad (forM_, when)
|
||||
import Control.Monad.IO.Class (liftIO)
|
||||
import Data.List (sortOn)
|
||||
import Data.Set qualified as Set
|
||||
import Data.UUID.V4 (nextRandom)
|
||||
import FL.Actors.Logger (logFn)
|
||||
import FL.CRDT
|
||||
import FL.NN (fedAvg)
|
||||
import FL.Types
|
||||
|
||||
data PeerCoordinatorState = PeerCoordinatorState
|
||||
{ pcsNodeId :: String
|
||||
, pcsRoundCounter :: GCounter -- completed training rounds per peer (monotone, never reset)
|
||||
, pcsParticipants :: ORSet String -- active peers ever seen (monotone, never reset)
|
||||
, pcsCurrentRound :: Int
|
||||
, pcsTotalRounds :: Int
|
||||
, pcsGlobalModel :: GlobalModel
|
||||
, pcsPeerRefs :: [ActorRef PeerCoordinatorMsg ()]
|
||||
, pcsSelfRef :: ActorRef PeerCoordinatorMsg ()
|
||||
, pcsTrainerRef :: ActorRef TrainerMsg ()
|
||||
, pcsEvalRef :: ActorRef EvaluatorMsg ()
|
||||
, pcsLogRef :: ActorRef LogMsg ()
|
||||
, pcsCollected :: [ModelUpdate] -- updates for the current round + buffered future rounds
|
||||
, pcsNumPeers :: Int -- total number of peers (including self)
|
||||
, pcsDoneVar :: MVar ()
|
||||
}
|
||||
|
||||
-- | Insert an update, deduplicating on (trainer, round). Stale (past-round)
|
||||
-- updates are dropped; future-round updates from fast peers are buffered so
|
||||
-- they are not lost — without the round key a fast peer's round-(r+1) update
|
||||
-- would overwrite its round-r update in a slower peer and rounds would mix.
|
||||
insertUpdate :: Int -> ModelUpdate -> [ModelUpdate] -> [ModelUpdate]
|
||||
insertUpdate currentRound upd collected
|
||||
| muRound upd < currentRound = collected
|
||||
| otherwise =
|
||||
upd : filter (\u -> (muTrainerId u, muRound u) /= (muTrainerId upd, muRound upd)) collected
|
||||
|
||||
currentRoundUpdates :: PeerCoordinatorState -> [ModelUpdate]
|
||||
currentRoundUpdates st = [u | u <- pcsCollected st, muRound u == pcsCurrentRound st]
|
||||
|
||||
-- | A round is complete when the CRDTs agree with the collected updates:
|
||||
-- the OR-Set has seen all peers, the G-Counter shows every one of them has
|
||||
-- finished training the current round, and we hold all their updates.
|
||||
roundComplete :: PeerCoordinatorState -> Bool
|
||||
roundComplete st =
|
||||
let r = pcsCurrentRound st
|
||||
members = orMembers (pcsParticipants st)
|
||||
trained = Set.filter (\p -> gcGet p (pcsRoundCounter st) > r) members
|
||||
in Set.size members >= pcsNumPeers st
|
||||
&& Set.size trained >= pcsNumPeers st
|
||||
&& length (currentRoundUpdates st) >= pcsNumPeers st
|
||||
|
||||
peerCoordinatorHandler :: Handler PeerCoordinatorMsg PeerCoordinatorState ()
|
||||
peerCoordinatorHandler (SetPeers peerIds) = do
|
||||
st <- state
|
||||
(SomeActorRef self) <- getSelf
|
||||
let peers = map (\aid -> RemoteRef aid :: ActorRef PeerCoordinatorMsg ()) peerIds
|
||||
selfRef = unsafeCoerce self :: ActorRef PeerCoordinatorMsg ()
|
||||
passWith st { pcsPeerRefs = peers, pcsSelfRef = selfRef }
|
||||
|
||||
peerCoordinatorHandler PeerStartRound = do
|
||||
st <- state
|
||||
logFn (pcsLogRef st) (pcsNodeId st) INFO $
|
||||
"Starting local training for P2P round " <> show (pcsCurrentRound st + 1)
|
||||
cast (StartTraining (pcsGlobalModel st)) (pcsTrainerRef st)
|
||||
pass
|
||||
|
||||
peerCoordinatorHandler (PeerTrainingDone upd) = do
|
||||
st <- state
|
||||
uuid <- liftIO nextRandom
|
||||
let counter' = gcIncrement (pcsNodeId st) (pcsRoundCounter st)
|
||||
participants' = orAdd uuid (pcsNodeId st) (pcsParticipants st)
|
||||
collected' = insertUpdate (pcsCurrentRound st) upd (pcsCollected st)
|
||||
payload = PeerSyncPayload
|
||||
{ pspCounter = counter'
|
||||
, pspParticipants = participants'
|
||||
, pspUpdate = Just upd
|
||||
}
|
||||
logFn (pcsLogRef st) (pcsNodeId st) INFO
|
||||
"Training done, broadcasting PeerSync"
|
||||
forM_ (pcsPeerRefs st) $ \ref ->
|
||||
cast (PeerSync payload) ref
|
||||
let st' = st { pcsRoundCounter = counter'
|
||||
, pcsParticipants = participants'
|
||||
, pcsCollected = collected'
|
||||
}
|
||||
if roundComplete st'
|
||||
then finishRound st'
|
||||
else passWith st'
|
||||
|
||||
peerCoordinatorHandler (PeerSync payload) = do
|
||||
st <- state
|
||||
-- Merge CRDTs
|
||||
let counter' = gcMerge (pcsRoundCounter st) (pspCounter payload)
|
||||
participants' = orMerge (pcsParticipants st) (pspParticipants payload)
|
||||
collected' = case pspUpdate payload of
|
||||
Nothing -> pcsCollected st
|
||||
Just upd -> insertUpdate (pcsCurrentRound st) upd (pcsCollected st)
|
||||
st' = st { pcsRoundCounter = counter'
|
||||
, pcsParticipants = participants'
|
||||
, pcsCollected = collected'
|
||||
}
|
||||
if roundComplete st'
|
||||
then finishRound st'
|
||||
else passWith st'
|
||||
|
||||
peerCoordinatorHandler (PeerEvaluationDone metrics) = do
|
||||
st <- state
|
||||
logFn (pcsLogRef st) (pcsNodeId st) INFO $
|
||||
"Eval round=" <> show (mRound metrics)
|
||||
<> " acc=" <> show4 (mAccuracy metrics)
|
||||
<> " auc=" <> show4 (mAUCROC metrics)
|
||||
liftIO $ when (mRound metrics >= pcsTotalRounds st) $
|
||||
putMVar (pcsDoneVar st) ()
|
||||
pass
|
||||
where
|
||||
show4 x = show (fromIntegral (round (x * 10000) :: Int) / 10000.0 :: Double)
|
||||
|
||||
finishRound :: PeerCoordinatorState -> Actor PeerCoordinatorMsg PeerCoordinatorState ()
|
||||
finishRound st = do
|
||||
-- Sort so every peer folds the weighted average in the same order and
|
||||
-- arrives at bit-identical models regardless of message arrival order.
|
||||
let updates = sortOn muTrainerId (currentRoundUpdates st)
|
||||
pairs = [(muSampleSize u, wData (muWeights u)) | u <- updates]
|
||||
avgW = Weights (fedAvg pairs)
|
||||
nextRound = pcsCurrentRound st + 1
|
||||
newModel = GlobalModel nextRound avgW
|
||||
logFn (pcsLogRef st) (pcsNodeId st) INFO $
|
||||
"P2P round " <> show (pcsCurrentRound st) <> " aggregated ("
|
||||
<> show (length updates) <> " updates)"
|
||||
cast (EvaluateModel newModel) (pcsEvalRef st)
|
||||
let st' = st
|
||||
{ pcsGlobalModel = newModel
|
||||
, pcsCurrentRound = nextRound
|
||||
-- Keep buffered updates from peers already in later rounds
|
||||
, pcsCollected = [u | u <- pcsCollected st, muRound u >= nextRound]
|
||||
}
|
||||
if nextRound >= pcsTotalRounds st
|
||||
then do
|
||||
let filename = "fl_model_" <> pcsNodeId st' <> ".bin"
|
||||
liftIO $ encodeFile filename (gmWeights newModel)
|
||||
logFn (pcsLogRef st') (pcsNodeId st') INFO $
|
||||
"Model saved to " <> filename
|
||||
passWith st'
|
||||
else do
|
||||
-- Schedule next round after a short delay to let PeerSync propagate
|
||||
castIn 500 PeerStartRound (pcsSelfRef st')
|
||||
passWith st'
|
||||
|
||||
peerCoordinatorDeathHandler :: DeathMessage -> ActorM PeerCoordinatorState (SupervisorAction PeerCoordinatorState)
|
||||
peerCoordinatorDeathHandler _ = do
|
||||
st <- state
|
||||
return (Continue st)
|
||||
@@ -0,0 +1,45 @@
|
||||
module FL.Actors.Trainer
|
||||
( TrainerState (..)
|
||||
, trainerHandler
|
||||
) where
|
||||
|
||||
import Control.Actor.Core
|
||||
import Control.Actor.Types
|
||||
import Control.Monad.IO.Class (liftIO)
|
||||
import FL.Actors.Logger (logFn)
|
||||
import FL.NN (LocalData, flattenWeights, localDataSize, trainLocal, unflattenWeights)
|
||||
import FL.Types
|
||||
|
||||
data TrainerState = TrainerState
|
||||
{ tsId :: String
|
||||
, tsData :: LocalData
|
||||
, tsLastRound :: Int
|
||||
, tsAggRef :: ActorRef AggregatorMsg ()
|
||||
, tsLogRef :: ActorRef LogMsg ()
|
||||
, tsConfig :: FLConfig
|
||||
}
|
||||
|
||||
trainerHandler :: Handler TrainerMsg TrainerState ()
|
||||
trainerHandler (StartTraining gm) = do
|
||||
st <- state
|
||||
let round_ = gmRound gm
|
||||
if round_ <= tsLastRound st
|
||||
then pass -- idempotent: already processed this round
|
||||
else do
|
||||
logFn (tsLogRef st) (tsId st) INFO $
|
||||
"Starting local training for round " <> show round_
|
||||
let initW = unflattenWeights (wData (gmWeights gm))
|
||||
updatedW <- liftIO $ trainLocal (tsConfig st) (tsId st) round_ initW (tsData st)
|
||||
let n = localDataSize (tsData st)
|
||||
update = ModelUpdate
|
||||
{ muTrainerId = tsId st
|
||||
, muRound = round_
|
||||
, muSampleSize = n
|
||||
, muWeights = Weights (flattenWeights updatedW)
|
||||
}
|
||||
logFn (tsLogRef st) (tsId st) INFO $
|
||||
"Round " <> show round_ <> " training complete, samples=" <> show n
|
||||
cast (SubmitUpdate update) (tsAggRef st)
|
||||
passWith st { tsLastRound = round_ }
|
||||
|
||||
trainerHandler TrainerStop = pass
|
||||
@@ -0,0 +1,82 @@
|
||||
module FL.CRDT
|
||||
( GCounter
|
||||
, gcEmpty
|
||||
, gcIncrement
|
||||
, gcMerge
|
||||
, gcValue
|
||||
, gcGet
|
||||
, ORSet
|
||||
, orEmpty
|
||||
, orAdd
|
||||
, orRemove
|
||||
, orMerge
|
||||
, orMembers
|
||||
) where
|
||||
|
||||
import Data.Binary (Binary (..))
|
||||
import Data.Map.Strict (Map)
|
||||
import Data.Map.Strict qualified as Map
|
||||
import Data.Set (Set)
|
||||
import Data.Set qualified as Set
|
||||
import Data.UUID (UUID)
|
||||
import GHC.Generics (Generic)
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- G-Counter: each node has its own non-decreasing counter.
|
||||
-- Merge takes element-wise max. Value is the sum.
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
newtype GCounter = GCounter (Map String Int)
|
||||
deriving (Show, Eq, Generic)
|
||||
|
||||
instance Binary GCounter where
|
||||
put (GCounter m) = put (Map.toAscList m)
|
||||
get = GCounter . Map.fromAscList <$> get
|
||||
|
||||
gcEmpty :: GCounter
|
||||
gcEmpty = GCounter Map.empty
|
||||
|
||||
gcIncrement :: String -> GCounter -> GCounter
|
||||
gcIncrement nodeId (GCounter m) =
|
||||
GCounter (Map.insertWith (+) nodeId 1 m)
|
||||
|
||||
gcMerge :: GCounter -> GCounter -> GCounter
|
||||
gcMerge (GCounter a) (GCounter b) =
|
||||
GCounter (Map.unionWith max a b)
|
||||
|
||||
gcValue :: GCounter -> Int
|
||||
gcValue (GCounter m) = sum (Map.elems m)
|
||||
|
||||
gcGet :: String -> GCounter -> Int
|
||||
gcGet nodeId (GCounter m) = Map.findWithDefault 0 nodeId m
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- OR-Set: observe-remove set. Each element tagged with a set of UUIDs.
|
||||
-- Add: insert a fresh UUID tag. Remove: delete all tags for that element.
|
||||
-- Merge: union of tag sets per element. An element is present iff it has
|
||||
-- at least one tag.
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
newtype ORSet a = ORSet (Map a (Set UUID))
|
||||
deriving (Show, Eq, Generic)
|
||||
|
||||
instance (Ord a, Binary a) => Binary (ORSet a) where
|
||||
put (ORSet m) = put (Map.toAscList (fmap Set.toAscList m))
|
||||
get = ORSet . Map.fromAscList . map (fmap Set.fromAscList) <$> get
|
||||
|
||||
orEmpty :: ORSet a
|
||||
orEmpty = ORSet Map.empty
|
||||
|
||||
orAdd :: Ord a => UUID -> a -> ORSet a -> ORSet a
|
||||
orAdd tag x (ORSet m) =
|
||||
ORSet (Map.insertWith Set.union x (Set.singleton tag) m)
|
||||
|
||||
orRemove :: Ord a => a -> ORSet a -> ORSet a
|
||||
orRemove x (ORSet m) = ORSet (Map.delete x m)
|
||||
|
||||
orMerge :: Ord a => ORSet a -> ORSet a -> ORSet a
|
||||
orMerge (ORSet a) (ORSet b) =
|
||||
ORSet (Map.unionWith Set.union a b)
|
||||
|
||||
orMembers :: Ord a => ORSet a -> Set a
|
||||
orMembers (ORSet m) = Map.keysSet (Map.filter (not . Set.null) m)
|
||||
@@ -0,0 +1,231 @@
|
||||
{-# LANGUAGE RecordWildCards #-}
|
||||
{-# LANGUAGE OverloadedStrings #-}
|
||||
|
||||
module FL.Data.PaySim
|
||||
( loadAndPartition
|
||||
-- , syntheticPartition
|
||||
) where
|
||||
|
||||
import Data.ByteString.Lazy qualified as BL
|
||||
import Data.Csv qualified as Csv
|
||||
import Data.List (sortBy)
|
||||
import Data.Ord (comparing)
|
||||
import Data.Vector qualified as V
|
||||
import FL.NN (LocalData (..))
|
||||
import FL.Types (FLConfig (..), NormMode (..), PartitionMode (..))
|
||||
import GHC.Generics (Generic)
|
||||
import Numeric.LinearAlgebra hiding ((<>))
|
||||
import Numeric.LinearAlgebra qualified as LA
|
||||
import System.Random (mkStdGen, randomRs)
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- PaySim CSV record
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
data PaySimRow = PaySimRow
|
||||
{ rowStep :: !Double
|
||||
, rowType :: !String
|
||||
, rowAmount :: !Double
|
||||
, rowOldBalOrig :: !Double
|
||||
, rowNewBalOrig :: !Double
|
||||
, rowOldBalDest :: !Double
|
||||
, rowNewBalDest :: !Double
|
||||
, rowIsFraud :: !Double
|
||||
} deriving (Show, Generic)
|
||||
|
||||
instance Csv.FromNamedRecord PaySimRow where
|
||||
parseNamedRecord r =
|
||||
PaySimRow
|
||||
<$> r Csv..: "step"
|
||||
<*> r Csv..: "type"
|
||||
<*> r Csv..: "amount"
|
||||
<*> r Csv..: "oldbalanceOrg"
|
||||
<*> r Csv..: "newbalanceOrig"
|
||||
<*> r Csv..: "oldbalanceDest"
|
||||
<*> r Csv..: "newbalanceDest"
|
||||
<*> r Csv..: "isFraud"
|
||||
|
||||
rowToFeatures :: PaySimRow -> Vector Double
|
||||
rowToFeatures PaySimRow{..} =
|
||||
let typeVec = case rowType of
|
||||
"CASH_IN" -> [1,0,0,0,0]
|
||||
"CASH_OUT" -> [0,1,0,0,0]
|
||||
"DEBIT" -> [0,0,1,0,0]
|
||||
"PAYMENT" -> [0,0,0,1,0]
|
||||
"TRANSFER" -> [0,0,0,0,1]
|
||||
_ -> [0,0,0,0,0]
|
||||
numeric = [ rowAmount
|
||||
, rowOldBalOrig
|
||||
, rowNewBalOrig
|
||||
, rowOldBalDest
|
||||
, rowNewBalDest
|
||||
, rowNewBalOrig - rowOldBalOrig -- balance change orig
|
||||
, rowNewBalDest - rowOldBalDest -- balance change dest
|
||||
, rowAmount / (rowOldBalOrig + 1) -- amount ratio orig
|
||||
, rowAmount / (rowOldBalDest + 1) -- amount ratio dest
|
||||
, rowStep / 744 -- normalized step (max 744 in PaySim)
|
||||
]
|
||||
-- pad to 29 features (5 one-hot + 10 numeric + 14 zeros)
|
||||
padded = typeVec ++ numeric ++ replicate 14 0.0
|
||||
in fromList padded
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Load PaySim CSV
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
loadPaySimCSV :: FilePath -> IO (Either String (Matrix Double, Vector Double))
|
||||
loadPaySimCSV path = do
|
||||
raw <- BL.readFile path
|
||||
case Csv.decodeByName raw of
|
||||
Left err -> return (Left err)
|
||||
Right (_, rows) -> do
|
||||
let featureList = V.toList (V.map rowToFeatures rows)
|
||||
labelList = V.toList (V.map rowIsFraud rows)
|
||||
return (Right (fromRows featureList, fromList labelList))
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Normalization
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
normalizeMinMax :: Matrix Double -> (Matrix Double, Vector Double, Vector Double)
|
||||
normalizeMinMax m =
|
||||
let mins = fromList [minElement (m ¿ [j]) | j <- [0 .. cols m - 1]]
|
||||
maxs = fromList [maxElement (m ¿ [j]) | j <- [0 .. cols m - 1]]
|
||||
range = cmap (\x -> if x == 0 then 1 else x) (maxs - mins)
|
||||
norm = (m - LA.repmat (asRow mins) (rows m) 1)
|
||||
/ LA.repmat (asRow range) (rows m) 1
|
||||
in (norm, mins, maxs)
|
||||
|
||||
normalizeZScore :: Matrix Double -> (Matrix Double, Vector Double, Vector Double)
|
||||
normalizeZScore m =
|
||||
let n = fromIntegral (rows m) :: Double
|
||||
means = fromList [sumElements (m ¿ [j]) / n | j <- [0 .. cols m - 1]]
|
||||
vars = fromList [sumElements (cmap (^(2::Int)) ((m ¿ [j]) - scalar (means ! j))) / n
|
||||
| j <- [0 .. cols m - 1]]
|
||||
stds = cmap (\v -> if v == 0 then 1 else sqrt v) vars
|
||||
norm = (m - LA.repmat (asRow means) (rows m) 1)
|
||||
/ LA.repmat (asRow stds) (rows m) 1
|
||||
in (norm, means, stds)
|
||||
|
||||
applyNorm :: NormMode -> (Matrix Double, Vector Double, Vector Double) -> Matrix Double -> Matrix Double
|
||||
applyNorm mode (_, p1, p2) m =
|
||||
case mode of
|
||||
MinMax ->
|
||||
let range = cmap (\x -> if x == 0 then 1 else x) (p2 - p1)
|
||||
in (m - LA.repmat (asRow p1) (rows m) 1) / LA.repmat (asRow range) (rows m) 1
|
||||
ZScore ->
|
||||
let stds = cmap (\x -> if x == 0 then 1 else x) p2
|
||||
in (m - LA.repmat (asRow p1) (rows m) 1) / LA.repmat (asRow stds) (rows m) 1
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- IID partition: deterministic shard by trainer index
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
partitionIID :: Int -> Int -> Matrix Double -> Vector Double -> LocalData
|
||||
partitionIID trainerIdx numTrainers feats labels =
|
||||
let n = rows feats
|
||||
-- Shuffle with fixed seed for reproducibility across processes
|
||||
seed = 2024
|
||||
rng = mkStdGen seed
|
||||
sorted = sortBy (comparing fst)
|
||||
(zip (randomRs (0.0 :: Double, 1.0) rng) [0 .. n - 1])
|
||||
idxs = map snd sorted
|
||||
shard = chunkedShard trainerIdx numTrainers idxs
|
||||
f = feats ? shard
|
||||
l = fromList [labels ! i | i <- shard]
|
||||
in LocalData f l
|
||||
|
||||
chunkedShard :: Int -> Int -> [a] -> [a]
|
||||
chunkedShard idx total xs =
|
||||
let n = length xs
|
||||
size = n `div` total
|
||||
start = idx * size
|
||||
end = if idx == total - 1 then n else start + size
|
||||
in take (end - start) (drop start xs)
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Non-IID partition: each trainer gets dominant transaction type
|
||||
-- (approximated via deterministic row assignment)
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
partitionNonIID :: Int -> Int -> Matrix Double -> Vector Double -> LocalData
|
||||
partitionNonIID trainerIdx numTrainers feats labels =
|
||||
let n = rows feats
|
||||
-- Assign rows to trainers based on dominant type one-hot (cols 0–4)
|
||||
-- Trainer i "owns" type (i mod 5); remaining trainers share type 0
|
||||
ownType = trainerIdx `mod` 5
|
||||
isOwn i = let row = toList (flatten (feats ? [i]))
|
||||
in length row > ownType && row !! ownType > 0.5
|
||||
-- With more than 5 trainers several trainers own the same type;
|
||||
-- split that type's rows among them instead of giving each trainer
|
||||
-- all of them (which duplicated data across shards).
|
||||
sharers = [i | i <- [0 .. numTrainers - 1], i `mod` 5 == ownType]
|
||||
shareIdx = length (takeWhile (/= trainerIdx) sharers)
|
||||
-- Gather own-type rows + 10% of other rows for minority
|
||||
ownIdxs = chunkedShard shareIdx (length sharers) [i | i <- [0 .. n-1], isOwn i]
|
||||
otherIdxs = [i | i <- [0 .. n-1], not (isOwn i)]
|
||||
minority = chunkedShard trainerIdx numTrainers otherIdxs
|
||||
allIdxs = ownIdxs ++ minority
|
||||
f = feats ? allIdxs
|
||||
l = fromList [labels ! i | i <- allIdxs]
|
||||
in LocalData f l
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Synthetic data fallback
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
syntheticData :: Int -> Int -> Int -> LocalData
|
||||
syntheticData trainerIdx numTrainers nSamples =
|
||||
let seed = trainerIdx * 13337 + numTrainers
|
||||
rng = mkStdGen seed
|
||||
-- 29 features from normal distribution
|
||||
featVals = take (nSamples * 29) (randomRs (-3.0 :: Double, 3.0) rng)
|
||||
feats = (nSamples >< 29) featVals
|
||||
-- Base 2% fraud rate, +1% per trainer index for variety (real PaySim
|
||||
-- is ~0.1%, but synthetic runs need enough positives to learn from)
|
||||
fraudRate = 0.02 + fromIntegral trainerIdx * 0.01
|
||||
rng2 = mkStdGen (seed + 1)
|
||||
labelVals = take nSamples (randomRs (0.0 :: Double, 1.0) rng2)
|
||||
labels = fromList [if v < fraudRate then 1.0 else 0.0 | v <- labelVals]
|
||||
in LocalData feats labels
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Top-level entry point used by Provider and P2P startup
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
-- Returns (train partition, test data).
|
||||
-- Test data is the first 10% of the full dataset (at least 1000 rows),
|
||||
-- taken before partitioning so every node evaluates on the same split.
|
||||
loadAndPartition
|
||||
:: FLConfig
|
||||
-> Int -- this node's trainer index (0-based)
|
||||
-> IO (LocalData, LocalData)
|
||||
loadAndPartition FLConfig{..} trainerIdx = do
|
||||
result <- case cfgDataPath of
|
||||
Nothing -> return (Left "no data path")
|
||||
Just path -> loadPaySimCSV path
|
||||
case result of
|
||||
Left _ -> do
|
||||
-- Synthetic fallback
|
||||
let testData = syntheticData (cfgNumTrainers + 1) cfgNumTrainers 2000
|
||||
trainData = syntheticData trainerIdx cfgNumTrainers 10000
|
||||
return (trainData, testData)
|
||||
Right (feats, labels) -> do
|
||||
let n = rows feats
|
||||
-- Aim for 1000 test rows for stable metrics, but never take more
|
||||
-- than half the data — max alone made testN exceed n on small
|
||||
-- CSVs and crashed the train subMatrix with a negative size.
|
||||
testN = min (max 1000 (n `div` 10)) (n `div` 2)
|
||||
testFeats = subMatrix (0, 0) (testN, cols feats) feats
|
||||
testLabs = subVector 0 testN labels
|
||||
trainF = subMatrix (testN, 0) (n - testN, cols feats) feats
|
||||
trainL = subVector testN (n - testN) labels
|
||||
(normF, p1, p2) = case cfgNorm of
|
||||
MinMax -> normalizeMinMax trainF
|
||||
ZScore -> normalizeZScore trainF
|
||||
normTest = applyNorm cfgNorm (normF, p1, p2) testFeats
|
||||
trainData = case cfgPartition of
|
||||
IID -> partitionIID trainerIdx cfgNumTrainers normF trainL
|
||||
NonIID -> partitionNonIID trainerIdx cfgNumTrainers normF trainL
|
||||
testData = LocalData normTest testLabs
|
||||
return (trainData, testData)
|
||||
+290
@@ -0,0 +1,290 @@
|
||||
{-# LANGUAGE RecordWildCards #-}
|
||||
{-# LANGUAGE StrictData #-}
|
||||
|
||||
module FL.NN
|
||||
( MLPWeights (..)
|
||||
, LocalData (..)
|
||||
, initWeights
|
||||
, trainLocal
|
||||
, evaluate
|
||||
, flattenWeights
|
||||
, unflattenWeights
|
||||
, fedAvg
|
||||
, localDataSize
|
||||
) where
|
||||
|
||||
import Data.List (foldl', sortBy)
|
||||
import Data.Ord (Down (..), comparing)
|
||||
import FL.Types (FLConfig (..), Metrics (..))
|
||||
import Numeric.LinearAlgebra hiding ((<>))
|
||||
import Numeric.LinearAlgebra qualified as LA
|
||||
import System.Random (mkStdGen, randomRs)
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Types
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
data MLPWeights = MLPWeights
|
||||
{ w1 :: Matrix Double -- 64 × 29
|
||||
, b1 :: Vector Double -- 64
|
||||
, w2 :: Matrix Double -- 32 × 64
|
||||
, b2 :: Vector Double -- 32
|
||||
, w3 :: Matrix Double -- 2 × 32
|
||||
, b3 :: Vector Double -- 2
|
||||
}
|
||||
|
||||
data LocalData = LocalData
|
||||
{ ldFeatures :: Matrix Double -- N × 29
|
||||
, ldLabels :: Vector Double -- N (0.0 or 1.0)
|
||||
}
|
||||
|
||||
localDataSize :: LocalData -> Int
|
||||
localDataSize = rows . ldFeatures
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Weight initialization (Xavier uniform)
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
xavierMatrix :: Int -> Int -> Int -> Matrix Double
|
||||
xavierMatrix seed fanIn fanOut =
|
||||
let bound = sqrt (6.0 / fromIntegral (fanIn + fanOut))
|
||||
n = fanIn * fanOut
|
||||
vals = take n (randomRs (-bound, bound) (mkStdGen seed))
|
||||
in (fanOut >< fanIn) vals
|
||||
|
||||
initWeights :: IO MLPWeights
|
||||
initWeights = do
|
||||
-- Use fixed seeds so multiple processes with same call order get different
|
||||
-- weight initializations only if we pass different seeds; for now deterministic.
|
||||
let w1' = xavierMatrix 42 29 64
|
||||
b1' = konst 0 64
|
||||
w2' = xavierMatrix 137 64 32
|
||||
b2' = konst 0 32
|
||||
w3' = xavierMatrix 999 32 2
|
||||
b3' = konst 0 2
|
||||
return MLPWeights{ w1=w1', b1=b1', w2=w2', b2=b2', w3=w3', b3=b3' }
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Activation functions
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
relu :: Matrix Double -> Matrix Double
|
||||
relu = cmap (max 0)
|
||||
|
||||
reluGrad :: Matrix Double -> Matrix Double
|
||||
reluGrad = cmap (\x -> if x > 0 then 1 else 0)
|
||||
|
||||
softmaxRows :: Matrix Double -> Matrix Double
|
||||
softmaxRows m = fromRows (map softmaxVec (toRows m))
|
||||
where
|
||||
-- Subtract the row max before exponentiating so large logits
|
||||
-- can't overflow exp and poison the weights with NaN.
|
||||
softmaxVec v =
|
||||
let mx = maxElement v
|
||||
e = cmap (\x -> exp (x - mx)) v
|
||||
s = sumElements e
|
||||
in cmap (/ s) e
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Forward pass
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
-- Returns (output probs N×2, z1, a1, z2, a2) for backprop.
|
||||
forwardWithCache
|
||||
:: MLPWeights
|
||||
-> Matrix Double -- N × 29
|
||||
-> (Matrix Double, Matrix Double, Matrix Double, Matrix Double, Matrix Double)
|
||||
forwardWithCache MLPWeights{..} x =
|
||||
let z1 = x LA.<> tr w1 + repRows (rows x) b1
|
||||
a1 = relu z1
|
||||
z2 = a1 LA.<> tr w2 + repRows (rows x) b2
|
||||
a2 = relu z2
|
||||
z3 = a2 LA.<> tr w3 + repRows (rows x) b3
|
||||
out = softmaxRows z3
|
||||
in (out, z1, a1, z2, a2)
|
||||
|
||||
repRows :: Int -> Vector Double -> Matrix Double
|
||||
repRows n v = fromRows (replicate n v)
|
||||
|
||||
colSums :: Matrix Double -> Vector Double
|
||||
colSums m = fromList [sumElements (m ¿ [j]) | j <- [0 .. cols m - 1]]
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Loss
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
crossEntropyLoss :: Matrix Double -> Vector Double -> Double
|
||||
crossEntropyLoss probs labels =
|
||||
let n = fromIntegral (rows probs)
|
||||
logProbs = cmap (log . max 1e-15) probs
|
||||
oneHot = buildOneHot (rows probs) (map round (toList labels))
|
||||
loss = negate (sumElements (oneHot * logProbs)) / n
|
||||
in loss
|
||||
|
||||
buildOneHot :: Int -> [Int] -> Matrix Double
|
||||
buildOneHot n ys =
|
||||
fromRows [fromList [if i == y then 1.0 else 0.0 | i <- [0, 1]] | y <- ys]
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Backward pass (mini-batch SGD)
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
stepBatch
|
||||
:: MLPWeights
|
||||
-> Matrix Double -- batch features N × 29
|
||||
-> Vector Double -- batch labels N
|
||||
-> Double -- learning rate
|
||||
-> (MLPWeights, Double)
|
||||
stepBatch ws@MLPWeights{..} x labels lr =
|
||||
let n = fromIntegral (rows x) :: Double
|
||||
(out, z1, a1, z2, a2) = forwardWithCache ws x
|
||||
oneHot = buildOneHot (rows x) (map round (toList labels))
|
||||
-- Output delta: dL/dZ3 = (probs - oneHot) / N
|
||||
dZ3 = cmap (/ n) (out - oneHot) -- N × 2
|
||||
dW3 = tr dZ3 LA.<> a2 -- 2 × 32
|
||||
db3 = colSums dZ3 -- 2
|
||||
dA2 = dZ3 LA.<> w3 -- N × 32
|
||||
dZ2 = dA2 * reluGrad z2 -- N × 32
|
||||
dW2 = tr dZ2 LA.<> a1 -- 32 × 64
|
||||
db2 = colSums dZ2 -- 32
|
||||
dA1 = dZ2 LA.<> w2 -- N × 64
|
||||
dZ1 = dA1 * reluGrad z1 -- N × 64
|
||||
dW1 = tr dZ1 LA.<> x -- 64 × 29
|
||||
db1 = colSums dZ1 -- 64
|
||||
ws' = MLPWeights
|
||||
{ w1 = w1 - cmap (* lr) dW1
|
||||
, b1 = b1 - cmap (* lr) db1
|
||||
, w2 = w2 - cmap (* lr) dW2
|
||||
, b2 = b2 - cmap (* lr) db2
|
||||
, w3 = w3 - cmap (* lr) dW3
|
||||
, b3 = b3 - cmap (* lr) db3
|
||||
}
|
||||
loss = crossEntropyLoss out labels
|
||||
in (ws', loss)
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Training
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
-- Shuffle row indices with a seed, then split into mini-batches.
|
||||
shuffleIndices :: Int -> Int -> [Int]
|
||||
shuffleIndices seed n =
|
||||
let rng = mkStdGen seed
|
||||
tagged = zip (randomRs (0.0 :: Double, 1.0) rng) [0 .. n - 1]
|
||||
in map snd (sortBy (comparing fst) tagged)
|
||||
|
||||
trainEpoch
|
||||
:: MLPWeights
|
||||
-> Matrix Double -- features N × 29
|
||||
-> Vector Double -- labels N
|
||||
-> Int -- batch size
|
||||
-> Double -- learning rate
|
||||
-> Int -- seed for shuffle
|
||||
-> (MLPWeights, Double)
|
||||
trainEpoch ws feats labels batchSize lr seed =
|
||||
let n = rows feats
|
||||
idxs = shuffleIndices seed n
|
||||
batches = chunksOf batchSize idxs
|
||||
step (w, totalLoss) batch =
|
||||
let bFeats = feats ? batch
|
||||
bLabels = fromList [labels ! i | i <- batch]
|
||||
(w', l) = stepBatch w bFeats bLabels lr
|
||||
in (w', totalLoss + l)
|
||||
(wFinal, lossSum) = foldl' step (ws, 0.0) batches
|
||||
in (wFinal, lossSum / fromIntegral (length batches))
|
||||
|
||||
chunksOf :: Int -> [a] -> [[a]]
|
||||
chunksOf _ [] = []
|
||||
chunksOf n xs = let (h, t) = splitAt n xs in h : chunksOf n t
|
||||
|
||||
-- | Train for the configured number of epochs. The shuffle seed is derived
|
||||
-- from the trainer id and round so batches differ across rounds and trainers
|
||||
-- (a bare epoch counter would replay the identical shuffle every round).
|
||||
trainLocal :: FLConfig -> String -> Int -> MLPWeights -> LocalData -> IO MLPWeights
|
||||
trainLocal FLConfig{..} trainerId roundNum initW LocalData{..} = do
|
||||
let seedBase = foldl' (\h c -> 31 * h + fromEnum c) (roundNum * 7919) trainerId
|
||||
go ws epoch =
|
||||
let (ws', _loss) = trainEpoch ws ldFeatures ldLabels cfgBatchSize cfgLearningRate (seedBase + epoch)
|
||||
in if epoch >= cfgEpochs then ws'
|
||||
else go ws' (epoch + 1)
|
||||
return (go initW 1)
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Evaluation
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
evaluate :: MLPWeights -> LocalData -> Int -> Metrics
|
||||
evaluate ws LocalData{..} round_ =
|
||||
let n = rows ldFeatures
|
||||
(probs, _, _, _, _) = forwardWithCache ws ldFeatures
|
||||
predLabels = fromList [if (probs ! i ! 1) >= 0.5 then 1.0 else 0.0 | i <- [0 .. n - 1]]
|
||||
scores = fromList [probs ! i ! 1 | i <- [0 .. n - 1]]
|
||||
tp = sumElements (predLabels * ldLabels)
|
||||
fp = sumElements (predLabels * (1 - ldLabels))
|
||||
fn = sumElements ((1 - predLabels) * ldLabels)
|
||||
tn = sumElements ((1 - predLabels) * (1 - ldLabels))
|
||||
acc = (tp + tn) / fromIntegral n
|
||||
prec = if tp + fp == 0 then 0 else tp / (tp + fp)
|
||||
rec = if tp + fn == 0 then 0 else tp / (tp + fn)
|
||||
f1 = if prec + rec == 0 then 0 else 2 * prec * rec / (prec + rec)
|
||||
auc = computeAUCROC (toList scores) (toList ldLabels)
|
||||
in Metrics acc prec rec f1 auc round_
|
||||
|
||||
computeAUCROC :: [Double] -> [Double] -> Double
|
||||
computeAUCROC scores labels =
|
||||
let nPos = fromIntegral (length (filter (== 1) labels))
|
||||
nNeg = fromIntegral (length (filter (== 0) labels))
|
||||
in if nPos == 0 || nNeg == 0
|
||||
then 0.5
|
||||
else
|
||||
let sorted = sortBy (comparing (Down . fst)) (zip scores labels)
|
||||
points = scanl (\(tp, fp) (_, y) -> (tp + y, fp + (1 - y))) (0, 0) sorted
|
||||
trapz (tp0, fp0) (tp1, fp1) =
|
||||
let fpr0 = fp0 / nNeg; tpr0 = tp0 / nPos
|
||||
fpr1 = fp1 / nNeg; tpr1 = tp1 / nPos
|
||||
in (fpr1 - fpr0) * (tpr0 + tpr1) / 2.0
|
||||
in sum (zipWith trapz points (tail points))
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Serialization bridge
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
flattenWeights :: MLPWeights -> [Double]
|
||||
flattenWeights MLPWeights{..} =
|
||||
toList (flatten w1)
|
||||
++ toList b1
|
||||
++ toList (flatten w2)
|
||||
++ toList b2
|
||||
++ toList (flatten w3)
|
||||
++ toList b3
|
||||
|
||||
-- Fixed architecture: 29→64→32→2
|
||||
-- w1: 64*29=1856, b1: 64, w2: 32*64=2048, b2: 32, w3: 2*32=64, b3: 2
|
||||
unflattenWeights :: [Double] -> MLPWeights
|
||||
unflattenWeights xs =
|
||||
let (xw1, r1) = splitAt (64*29) xs
|
||||
(xb1, r2) = splitAt 64 r1
|
||||
(xw2, r3) = splitAt (32*64) r2
|
||||
(xb2, r4) = splitAt 32 r3
|
||||
(xw3, r5) = splitAt (2*32) r4
|
||||
xb3 = take 2 r5
|
||||
in MLPWeights
|
||||
{ w1 = (64 >< 29) xw1
|
||||
, b1 = fromList xb1
|
||||
, w2 = (32 >< 64) xw2
|
||||
, b2 = fromList xb2
|
||||
, w3 = ( 2 >< 32) xw3
|
||||
, b3 = fromList xb3
|
||||
}
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- FedAvg
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
-- Weighted average of flat weight vectors.
|
||||
fedAvg :: [(Int, [Double])] -> [Double]
|
||||
fedAvg [] = []
|
||||
fedAvg updates =
|
||||
let totalN = fromIntegral (sum (map fst updates)) :: Double
|
||||
weighted (n, ws) = map (* (fromIntegral n / totalN)) ws
|
||||
in foldl1 (zipWith (+)) (map weighted updates)
|
||||
+182
@@ -0,0 +1,182 @@
|
||||
{-# LANGUAGE RecordWildCards #-}
|
||||
module FL.P2P
|
||||
( runPeer
|
||||
) where
|
||||
|
||||
import Control.Actor
|
||||
import Control.Actor.Registry (registerActor, lookupRemoteActor)
|
||||
import Control.Concurrent (threadDelay)
|
||||
import Control.Concurrent.MVar (newEmptyMVar, takeMVar)
|
||||
import Control.Concurrent.STM
|
||||
( TMVar, atomically, newEmptyTMVarIO, putTMVar, readTMVar, tryTakeTMVar )
|
||||
import Control.Monad (forM, forM_, void)
|
||||
import Control.Monad.IO.Class (liftIO)
|
||||
import FL.Actors.Logger (logFn, logFn', loggerHandler)
|
||||
import FL.Actors.PeerCoordinator
|
||||
import FL.CRDT (gcEmpty, orEmpty)
|
||||
import FL.Data.PaySim (loadAndPartition)
|
||||
import FL.NN
|
||||
( LocalData, evaluate, flattenWeights, initWeights
|
||||
, localDataSize, trainLocal, unflattenWeights
|
||||
)
|
||||
import FL.Retry (retryConnect, retryLookup)
|
||||
import FL.Types
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- P2P-specific actor state types
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
-- Evaluator reads the PeerCoordinator ref lazily from a TMVar.
|
||||
-- This breaks the circular dependency: evaluator is spawned before PC,
|
||||
-- but PC's ref is filled into the TMVar right after PC spawns.
|
||||
data P2PEvalState = P2PEvalState
|
||||
{ p2eTestData :: LocalData
|
||||
, p2ePcCell :: TMVar (ActorRef PeerCoordinatorMsg ())
|
||||
, p2eLogRef :: ActorRef LogMsg ()
|
||||
}
|
||||
|
||||
p2pEvalHandler :: Handler EvaluatorMsg P2PEvalState ()
|
||||
p2pEvalHandler (EvaluateModel gm) = do
|
||||
st <- state
|
||||
let ws = unflattenWeights (wData (gmWeights gm))
|
||||
metrics = evaluate ws (p2eTestData st) (gmRound gm)
|
||||
logFn (p2eLogRef st) "Evaluator" INFO $
|
||||
"P2P round " <> show (gmRound gm)
|
||||
<> " acc=" <> show4 (mAccuracy metrics)
|
||||
<> " auc=" <> show4 (mAUCROC metrics)
|
||||
pcRef <- liftIO $ atomically $ readTMVar (p2ePcCell st)
|
||||
cast (PeerEvaluationDone metrics) pcRef
|
||||
pass
|
||||
where
|
||||
show4 x = show (fromIntegral (round (x * 10000) :: Int) / 10000.0 :: Double)
|
||||
|
||||
-- Trainer likewise reads the PeerCoordinator ref lazily.
|
||||
data P2PTrainerState = P2PTrainerState
|
||||
{ p2tId :: String
|
||||
, p2tData :: LocalData
|
||||
, p2tLastRound :: Int
|
||||
, p2tPcCell :: TMVar (ActorRef PeerCoordinatorMsg ())
|
||||
, p2tLogRef :: ActorRef LogMsg ()
|
||||
, p2tConfig :: FLConfig
|
||||
}
|
||||
|
||||
p2pTrainerHandler :: Handler TrainerMsg P2PTrainerState ()
|
||||
p2pTrainerHandler (StartTraining gm) = do
|
||||
st <- state
|
||||
let round_ = gmRound gm
|
||||
if round_ <= p2tLastRound st
|
||||
then pass
|
||||
else do
|
||||
logFn (p2tLogRef st) (p2tId st) INFO $
|
||||
"P2P training round " <> show round_
|
||||
let initW = unflattenWeights (wData (gmWeights gm))
|
||||
updatedW <- liftIO $ trainLocal (p2tConfig st) (p2tId st) round_ initW (p2tData st)
|
||||
let n = localDataSize (p2tData st)
|
||||
upd = ModelUpdate (p2tId st) round_ n (Weights (flattenWeights updatedW))
|
||||
pcRef <- liftIO $ atomically $ readTMVar (p2tPcCell st)
|
||||
cast (PeerTrainingDone upd) pcRef
|
||||
passWith st { p2tLastRound = round_ }
|
||||
p2pTrainerHandler TrainerStop = pass
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Peer node startup
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
runPeer :: FLConfig -> Int -> [NodeAddr] -> IO ()
|
||||
runPeer cfg peerIdx knownPeers = do
|
||||
let FLConfig{..} = cfg
|
||||
peerPort = cfgBasePort
|
||||
peerId = "peer-" <> show peerIdx
|
||||
|
||||
rt <- initRuntime (NodeAddr "localhost" (fromIntegral peerPort)) createTCPTransport
|
||||
doneVar <- newEmptyMVar
|
||||
|
||||
withRuntime rt $ do
|
||||
forM_ knownPeers $ \addr -> retryConnect 20 addr
|
||||
liftIO $ threadDelay 300_000
|
||||
|
||||
logRef <- spawnActor loggerHandler stopOnDeath ()
|
||||
registerActor (peerId <> "-logger") logRef
|
||||
|
||||
initW <- liftIO initWeights
|
||||
let initModel = GlobalModel 0 (Weights (flattenWeights initW))
|
||||
|
||||
(trainData, testData) <- liftIO $ loadAndPartition cfg peerIdx
|
||||
|
||||
-- Shared TMVar: filled with the PeerCoordinator ref after it spawns.
|
||||
-- Evaluator and Trainer read from it lazily on their first message.
|
||||
pcCell <- liftIO newEmptyTMVarIO
|
||||
|
||||
evalRef <- spawnActor p2pEvalHandler stopOnDeath
|
||||
(P2PEvalState testData pcCell logRef)
|
||||
|
||||
trainerRef <- spawnActor p2pTrainerHandler stopOnDeath
|
||||
P2PTrainerState
|
||||
{ p2tId = peerId
|
||||
, p2tData = trainData
|
||||
, p2tLastRound = -1
|
||||
, p2tPcCell = pcCell
|
||||
, p2tLogRef = logRef
|
||||
, p2tConfig = cfg
|
||||
}
|
||||
|
||||
-- PeerCoordinator — placeholder self-ref (replaced by SetPeers)
|
||||
let dummyId = ActorId 0 (read "00000000-0000-0000-0000-000000000002")
|
||||
dummySelf = RemoteRef dummyId :: ActorRef PeerCoordinatorMsg ()
|
||||
|
||||
let pcSt = PeerCoordinatorState
|
||||
{ pcsNodeId = peerId
|
||||
, pcsRoundCounter = gcEmpty
|
||||
, pcsParticipants = orEmpty
|
||||
, pcsCurrentRound = 0
|
||||
, pcsTotalRounds = cfgRounds
|
||||
, pcsGlobalModel = initModel
|
||||
, pcsPeerRefs = []
|
||||
, pcsSelfRef = dummySelf
|
||||
, pcsTrainerRef = trainerRef
|
||||
, pcsEvalRef = evalRef
|
||||
, pcsLogRef = logRef
|
||||
, pcsCollected = []
|
||||
, pcsNumPeers = cfgNumTrainers
|
||||
, pcsDoneVar = doneVar
|
||||
}
|
||||
|
||||
-- Supervised PeerCoordinator. On every (re)spawn the shared cell is
|
||||
-- refreshed and the registry name re-registered, so trainer/evaluator
|
||||
-- keep reaching the live incarnation. (State restarts fresh.)
|
||||
supervise' OneForOne
|
||||
[ ChildSpec
|
||||
{ csRun = \target -> do
|
||||
ref <- spawnActor peerCoordinatorHandler peerCoordinatorDeathHandler pcSt
|
||||
linkActorTo target ref
|
||||
return ref
|
||||
, csOnSpawn = \ref -> do
|
||||
liftIO $ atomically $ do
|
||||
void $ tryTakeTMVar pcCell
|
||||
putTMVar pcCell ref
|
||||
registerActor ("peercoord-" <> peerId) ref
|
||||
}
|
||||
]
|
||||
pcRef <- liftIO $ atomically $ readTMVar pcCell
|
||||
|
||||
logFn' logRef peerId INFO $
|
||||
"Peer ready, waiting for " <> show (cfgNumTrainers - 1) <> " other peers"
|
||||
|
||||
-- Discover all other peer coordinators
|
||||
let otherIds = ["peer-" <> show i | i <- [0 .. cfgNumTrainers - 1], i /= peerIdx]
|
||||
peerPcRefs <- forM otherIds $ \pid ->
|
||||
retryLookup 40
|
||||
(lookupRemoteActor ("peercoord-" <> pid)
|
||||
:: RuntimeM (Maybe (ActorRef PeerCoordinatorMsg ())))
|
||||
|
||||
-- Wire up PeerCoordinator with peers
|
||||
let peerIds = map actorRefId peerPcRefs
|
||||
cast' (SetPeers peerIds) pcRef
|
||||
|
||||
liftIO $ threadDelay 200_000
|
||||
cast' PeerStartRound pcRef
|
||||
logFn' logRef peerId INFO "P2P round 1 started"
|
||||
|
||||
takeMVar doneVar
|
||||
threadDelay 300_000
|
||||
putStrLn $ "\n=== Peer " <> show peerIdx <> " federation complete ==="
|
||||
@@ -0,0 +1,149 @@
|
||||
{-# LANGUAGE RecordWildCards #-}
|
||||
module FL.Provider
|
||||
( runCoordinator
|
||||
, runTrainer
|
||||
) where
|
||||
|
||||
import Control.Actor
|
||||
import Control.Actor.Registry (lookupRemoteActor, registerActor)
|
||||
import Control.Concurrent (threadDelay)
|
||||
import Control.Concurrent.MVar (newEmptyMVar, takeMVar)
|
||||
import Control.Concurrent.STM
|
||||
( atomically, newEmptyTMVarIO, putTMVar, readTMVar, tryTakeTMVar )
|
||||
import Control.Monad (void)
|
||||
import Control.Monad.IO.Class (liftIO)
|
||||
import Data.Map.Strict qualified as Map
|
||||
import FL.Actors.Aggregator
|
||||
import FL.Actors.Coordinator
|
||||
import FL.Actors.Evaluator
|
||||
import FL.Actors.Logger
|
||||
import FL.Actors.Trainer
|
||||
import FL.CRDT (orEmpty)
|
||||
import FL.Data.PaySim (loadAndPartition)
|
||||
import FL.NN (flattenWeights, initWeights)
|
||||
import FL.Retry (retryConnect, retryLookup)
|
||||
import FL.Types
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Coordinator node
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
runCoordinator :: FLConfig -> IO ()
|
||||
runCoordinator cfg = do
|
||||
let FLConfig{..} = cfg
|
||||
rt <- initRuntime (NodeAddr "localhost" (fromIntegral cfgBasePort)) createTCPTransport
|
||||
doneVar <- newEmptyMVar
|
||||
|
||||
withRuntime rt $ do
|
||||
-- Logger
|
||||
logRef <- spawnActor loggerHandler stopOnDeath ()
|
||||
registerActor "logger" logRef
|
||||
|
||||
-- Initial global model
|
||||
initW <- liftIO initWeights
|
||||
let initModel = GlobalModel 0 (Weights (flattenWeights initW))
|
||||
|
||||
-- Shared TMVar for coordinator ref: filled after coordinator spawns.
|
||||
-- Aggregator and Evaluator read from it lazily.
|
||||
coordCell <- liftIO newEmptyTMVarIO
|
||||
|
||||
-- Load test data (coordinator holds the shared test split)
|
||||
(_, testData) <- liftIO $ loadAndPartition cfg 0
|
||||
|
||||
-- Evaluator
|
||||
evalRef <- spawnActor evaluatorHandler stopOnDeath
|
||||
(EvaluatorState testData coordCell logRef)
|
||||
registerActor "evaluator" evalRef
|
||||
|
||||
-- Aggregator
|
||||
aggRef <- spawnActor aggregatorHandler stopOnDeath
|
||||
(AggregatorState cfgNumTrainers [] coordCell logRef 0)
|
||||
registerActor "aggregator" aggRef
|
||||
|
||||
-- Coordinator (knows aggRef and evalRef at spawn time)
|
||||
let initCoordSt = CoordinatorState
|
||||
{ csTrainers = orEmpty
|
||||
, csTrainerRefs = Map.empty
|
||||
, csExpectedTrainers = cfgNumTrainers
|
||||
, csCurrentRound = 0
|
||||
, csTotalRounds = cfgRounds
|
||||
, csGlobalModel = initModel
|
||||
, csAggRef = aggRef
|
||||
, csEvalRef = evalRef
|
||||
, csLogRef = logRef
|
||||
, csDoneVar = doneVar
|
||||
}
|
||||
-- Supervised coordinator. On every (re)spawn — including restarts — the
|
||||
-- shared cell is refreshed and the registry name re-registered, so
|
||||
-- aggregator/evaluator/trainers keep reaching the live incarnation.
|
||||
-- (Coordinator state itself restarts fresh: round progress is lost.)
|
||||
supervise' OneForOne
|
||||
[ ChildSpec
|
||||
{ csRun = \target -> do
|
||||
ref <- spawnActor coordinatorHandler coordinatorDeathHandler initCoordSt
|
||||
linkActorTo target ref
|
||||
return ref
|
||||
, csOnSpawn = \ref -> do
|
||||
liftIO $ atomically $ do
|
||||
void $ tryTakeTMVar coordCell
|
||||
putTMVar coordCell ref
|
||||
registerActor "coordinator" ref
|
||||
}
|
||||
]
|
||||
void $ liftIO $ atomically $ readTMVar coordCell
|
||||
|
||||
logFn' logRef "System" INFO $
|
||||
"Coordinator ready, waiting for " <> show cfgNumTrainers <> " trainers"
|
||||
|
||||
takeMVar doneVar
|
||||
threadDelay 300_000
|
||||
putStrLn "\n=== Provider federation complete ==="
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Trainer node
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
runTrainer :: FLConfig -> Int -> NodeAddr -> IO ()
|
||||
runTrainer cfg trainerIdx coordAddr = do
|
||||
let FLConfig{..} = cfg
|
||||
let trainerPort = cfgBasePort + trainerIdx + 1
|
||||
trainerId = "trainer-" <> show trainerIdx
|
||||
|
||||
rt <- initRuntime (NodeAddr "localhost" (fromIntegral trainerPort)) createTCPTransport
|
||||
|
||||
withRuntime rt $ do
|
||||
retryConnect 20 coordAddr
|
||||
liftIO $ threadDelay 500_000 -- wait for registry sync
|
||||
|
||||
-- Look up aggregator and logger on the coordinator node (retry until
|
||||
-- available; the logger is registered first, so if the aggregator
|
||||
-- resolves the logger will too)
|
||||
aggRef <- retryLookup 20
|
||||
(lookupRemoteActor "aggregator" :: RuntimeM (Maybe (ActorRef AggregatorMsg ())))
|
||||
logRef <- retryLookup 20
|
||||
(lookupRemoteActor "logger" :: RuntimeM (Maybe (ActorRef LogMsg ())))
|
||||
|
||||
(trainData, _) <- liftIO $ loadAndPartition cfg trainerIdx
|
||||
|
||||
let trainerSt = TrainerState
|
||||
{ tsId = trainerId
|
||||
, tsData = trainData
|
||||
, tsLastRound = -1
|
||||
, tsAggRef = aggRef
|
||||
, tsLogRef = logRef
|
||||
, tsConfig = cfg
|
||||
}
|
||||
trainerRef <- spawnActor trainerHandler stopOnDeath trainerSt
|
||||
registerActor trainerId trainerRef
|
||||
|
||||
-- Register with coordinator
|
||||
coordRef <- retryLookup 20
|
||||
(lookupRemoteActor "coordinator" :: RuntimeM (Maybe (ActorRef CoordinatorMsg ())))
|
||||
cast' (RegisterTrainer trainerId (actorRefId trainerRef)) coordRef
|
||||
|
||||
logFn' logRef trainerId INFO "Registered with coordinator"
|
||||
|
||||
waitForever
|
||||
|
||||
waitForever :: IO ()
|
||||
waitForever = threadDelay maxBound >> waitForever
|
||||
@@ -0,0 +1,34 @@
|
||||
-- | Startup helpers shared by the provider and P2P entry points: retry a
|
||||
-- node connection or a registry lookup until it succeeds or attempts run out.
|
||||
module FL.Retry
|
||||
( retryConnect
|
||||
, retryLookup
|
||||
) where
|
||||
|
||||
import Control.Actor (ActorRef, NodeAddr, RuntimeM, connect, withRuntime)
|
||||
import Control.Concurrent (threadDelay)
|
||||
import Control.Exception (SomeException, try)
|
||||
import Control.Monad (void)
|
||||
import Control.Monad.IO.Class (liftIO)
|
||||
import Control.Monad.Reader (ask)
|
||||
|
||||
retryConnect :: Int -> NodeAddr -> RuntimeM ()
|
||||
retryConnect 0 addr = error $ "retryConnect: could not connect to " <> show addr
|
||||
retryConnect n addr = do
|
||||
rt <- ask
|
||||
result <- liftIO $ try @SomeException $ withRuntime rt $ void (connect Nothing addr)
|
||||
case result of
|
||||
Right () -> return ()
|
||||
Left _ -> do
|
||||
liftIO $ threadDelay 1_000_000
|
||||
retryConnect (n - 1) addr
|
||||
|
||||
retryLookup :: Int -> RuntimeM (Maybe (ActorRef msg r)) -> RuntimeM (ActorRef msg r)
|
||||
retryLookup 0 _ = error "retryLookup: exhausted retries"
|
||||
retryLookup n action = do
|
||||
result <- action
|
||||
case result of
|
||||
Just ref -> return ref
|
||||
Nothing -> do
|
||||
liftIO $ threadDelay 500_000
|
||||
retryLookup (n - 1) action
|
||||
+160
@@ -0,0 +1,160 @@
|
||||
module FL.Types
|
||||
( -- Config
|
||||
FLConfig (..)
|
||||
, PartitionMode (..)
|
||||
, NormMode (..)
|
||||
-- Wire types
|
||||
, Weights (..)
|
||||
, GlobalModel (..)
|
||||
, ModelUpdate (..)
|
||||
, Metrics (..)
|
||||
-- Logger messages
|
||||
, LogLevel (..)
|
||||
, LogEntry (..)
|
||||
, LogMsg (..)
|
||||
-- Actor messages
|
||||
, EvaluatorMsg (..)
|
||||
, TrainerMsg (..)
|
||||
, AggregatorMsg (..)
|
||||
, CoordinatorMsg (..)
|
||||
, PeerSyncPayload (..)
|
||||
, PeerCoordinatorMsg (..)
|
||||
) where
|
||||
|
||||
import Control.Actor.Types (ActorId)
|
||||
import Data.Binary (Binary)
|
||||
import FL.CRDT (GCounter, ORSet)
|
||||
import GHC.Generics (Generic)
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Config (not a wire type — parsed from CLI, passed around locally)
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
data PartitionMode = IID | NonIID
|
||||
deriving (Show, Eq, Generic)
|
||||
instance Binary PartitionMode
|
||||
|
||||
data NormMode = MinMax | ZScore
|
||||
deriving (Show, Eq, Generic)
|
||||
instance Binary NormMode
|
||||
|
||||
data FLConfig = FLConfig
|
||||
{ cfgRounds :: Int
|
||||
, cfgEpochs :: Int
|
||||
, cfgLearningRate :: Double
|
||||
, cfgBatchSize :: Int
|
||||
, cfgNumTrainers :: Int
|
||||
, cfgDataPath :: Maybe FilePath
|
||||
, cfgPartition :: PartitionMode
|
||||
, cfgNorm :: NormMode
|
||||
, cfgBasePort :: Int
|
||||
} deriving (Show, Generic)
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Wire types (sent over the network, must be Binary)
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
newtype Weights = Weights { wData :: [Double] }
|
||||
deriving (Show, Eq, Generic)
|
||||
instance Binary Weights
|
||||
|
||||
data GlobalModel = GlobalModel
|
||||
{ gmRound :: Int
|
||||
, gmWeights :: Weights
|
||||
} deriving (Show, Eq, Generic)
|
||||
instance Binary GlobalModel
|
||||
|
||||
data ModelUpdate = ModelUpdate
|
||||
{ muTrainerId :: String
|
||||
, muRound :: Int
|
||||
, muSampleSize :: Int
|
||||
, muWeights :: Weights
|
||||
} deriving (Show, Eq, Generic)
|
||||
instance Binary ModelUpdate
|
||||
|
||||
data Metrics = Metrics
|
||||
{ mAccuracy :: Double
|
||||
, mPrecision :: Double
|
||||
, mRecall :: Double
|
||||
, mF1 :: Double
|
||||
, mAUCROC :: Double
|
||||
, mRound :: Int
|
||||
} deriving (Show, Eq, Generic)
|
||||
instance Binary Metrics
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Logger
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
data LogLevel = DEBUG | INFO | WARN | ERROR
|
||||
deriving (Show, Eq, Ord, Generic)
|
||||
instance Binary LogLevel
|
||||
|
||||
data LogEntry = LogEntry
|
||||
{ leSource :: String
|
||||
, leLevel :: LogLevel
|
||||
, leMsg :: String
|
||||
} deriving (Show, Generic)
|
||||
instance Binary LogEntry
|
||||
|
||||
data LogMsg = Log LogEntry
|
||||
deriving (Show, Generic)
|
||||
instance Binary LogMsg
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Evaluator
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
data EvaluatorMsg = EvaluateModel GlobalModel
|
||||
deriving (Show, Generic)
|
||||
instance Binary EvaluatorMsg
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Trainer
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
data TrainerMsg
|
||||
= StartTraining GlobalModel
|
||||
| TrainerStop
|
||||
deriving (Show, Generic)
|
||||
instance Binary TrainerMsg
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Aggregator (provider mode)
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
data AggregatorMsg = SubmitUpdate ModelUpdate
|
||||
deriving (Show, Generic)
|
||||
instance Binary AggregatorMsg
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- Coordinator (provider mode)
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
data CoordinatorMsg
|
||||
= RegisterTrainer String ActorId
|
||||
| AggregationComplete GlobalModel
|
||||
| EvaluationDone Metrics
|
||||
| StartFederation
|
||||
deriving (Show, Generic)
|
||||
instance Binary CoordinatorMsg
|
||||
|
||||
-- ---------------------------------------------------------------------------
|
||||
-- PeerCoordinator (P2P mode)
|
||||
-- ---------------------------------------------------------------------------
|
||||
|
||||
data PeerSyncPayload = PeerSyncPayload
|
||||
{ pspCounter :: !GCounter
|
||||
, pspParticipants :: !(ORSet String)
|
||||
, pspUpdate :: !(Maybe ModelUpdate)
|
||||
} deriving (Show, Generic)
|
||||
instance Binary PeerSyncPayload
|
||||
|
||||
data PeerCoordinatorMsg
|
||||
= SetPeers ![ActorId]
|
||||
| PeerStartRound
|
||||
| PeerTrainingDone !ModelUpdate
|
||||
| PeerSync !PeerSyncPayload
|
||||
| PeerEvaluationDone !Metrics
|
||||
deriving (Show, Generic)
|
||||
instance Binary PeerCoordinatorMsg
|
||||
Reference in New Issue
Block a user