From 711d26d9f760925ee9150a88585f737434ccbbad Mon Sep 17 00:00:00 2001 From: dvdrw Date: Sun, 21 Jun 2026 12:00:00 +0200 Subject: [PATCH] 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 --- .gitignore | 2 + actors.cabal | 39 ++++- app/Main.hs | 150 ++++++++++++++++ scripts/demo-p2p.sh | 85 +++++++++ scripts/demo-provider.sh | 72 ++++++++ src/Control/Actor.hs | 2 + src/Control/Actor/Core.hs | 42 ++++- src/Control/Actor/Network.hs | 188 ++++++++++++-------- src/Control/Actor/Registry.hs | 18 +- src/Control/Actor/Runtime.hs | 3 +- src/Control/Actor/Supervision.hs | 72 +++++--- src/Control/Actor/Types.hs | 11 +- src/FL/Actors/Aggregator.hs | 43 +++++ src/FL/Actors/Coordinator.hs | 103 +++++++++++ src/FL/Actors/Evaluator.hs | 36 ++++ src/FL/Actors/Logger.hs | 37 ++++ src/FL/Actors/PeerCoordinator.hs | 166 ++++++++++++++++++ src/FL/Actors/Trainer.hs | 45 +++++ src/FL/CRDT.hs | 82 +++++++++ src/FL/Data/PaySim.hs | 231 ++++++++++++++++++++++++ src/FL/NN.hs | 290 +++++++++++++++++++++++++++++++ src/FL/P2P.hs | 182 +++++++++++++++++++ src/FL/Provider.hs | 149 ++++++++++++++++ src/FL/Retry.hs | 34 ++++ src/FL/Types.hs | 160 +++++++++++++++++ 25 files changed, 2127 insertions(+), 115 deletions(-) create mode 100644 app/Main.hs create mode 100755 scripts/demo-p2p.sh create mode 100755 scripts/demo-provider.sh create mode 100644 src/FL/Actors/Aggregator.hs create mode 100644 src/FL/Actors/Coordinator.hs create mode 100644 src/FL/Actors/Evaluator.hs create mode 100644 src/FL/Actors/Logger.hs create mode 100644 src/FL/Actors/PeerCoordinator.hs create mode 100644 src/FL/Actors/Trainer.hs create mode 100644 src/FL/CRDT.hs create mode 100644 src/FL/Data/PaySim.hs create mode 100644 src/FL/NN.hs create mode 100644 src/FL/P2P.hs create mode 100644 src/FL/Provider.hs create mode 100644 src/FL/Retry.hs create mode 100644 src/FL/Types.hs diff --git a/.gitignore b/.gitignore index 48a004c..0449a6b 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,3 @@ dist-newstyle +demo-out/ +fl_model_*.bin diff --git a/actors.cabal b/actors.cabal index da6d888..748410d 100644 --- a/actors.cabal +++ b/actors.cabal @@ -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 diff --git a/app/Main.hs b/app/Main.hs new file mode 100644 index 0000000..88415df --- /dev/null +++ b/app/Main.hs @@ -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" diff --git a/scripts/demo-p2p.sh b/scripts/demo-p2p.sh new file mode 100755 index 0000000..b74e509 --- /dev/null +++ b/scripts/demo-p2p.sh @@ -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/" diff --git a/scripts/demo-provider.sh b/scripts/demo-provider.sh new file mode 100755 index 0000000..14e43b3 --- /dev/null +++ b/scripts/demo-provider.sh @@ -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 diff --git a/src/Control/Actor.hs b/src/Control/Actor.hs index 8e8c268..dae7e07 100644 --- a/src/Control/Actor.hs +++ b/src/Control/Actor.hs @@ -5,6 +5,8 @@ module Control.Actor , module Control.Actor.Core , module Control.Actor.Supervision , module Control.Actor.Network + -- Initialization + , initRuntime -- Demo , pingActor , forwardActorWithCell diff --git a/src/Control/Actor/Core.hs b/src/Control/Actor/Core.hs index 02b77f8..42f7d57 100644 --- a/src/Control/Actor/Core.hs +++ b/src/Control/Actor/Core.hs @@ -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' diff --git a/src/Control/Actor/Network.hs b/src/Control/Actor/Network.hs index de3ee19..f5a67a7 100644 --- a/src/Control/Actor/Network.hs +++ b/src/Control/Actor/Network.hs @@ -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 diff --git a/src/Control/Actor/Registry.hs b/src/Control/Actor/Registry.hs index c1f0bdb..a432393 100644 --- a/src/Control/Actor/Registry.hs +++ b/src/Control/Actor/Registry.hs @@ -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)) diff --git a/src/Control/Actor/Runtime.hs b/src/Control/Actor/Runtime.hs index cbc39e7..48e3300 100644 --- a/src/Control/Actor/Runtime.hs +++ b/src/Control/Actor/Runtime.hs @@ -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 () } diff --git a/src/Control/Actor/Supervision.hs b/src/Control/Actor/Supervision.hs index e7301ef..5779d99 100644 --- a/src/Control/Actor/Supervision.hs +++ b/src/Control/Actor/Supervision.hs @@ -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) diff --git a/src/Control/Actor/Types.hs b/src/Control/Actor/Types.hs index cc043a1..c8d14b2 100644 --- a/src/Control/Actor/Types.hs +++ b/src/Control/Actor/Types.hs @@ -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 diff --git a/src/FL/Actors/Aggregator.hs b/src/FL/Actors/Aggregator.hs new file mode 100644 index 0000000..69fdea8 --- /dev/null +++ b/src/FL/Actors/Aggregator.hs @@ -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' } diff --git a/src/FL/Actors/Coordinator.hs b/src/FL/Actors/Coordinator.hs new file mode 100644 index 0000000..b8b6861 --- /dev/null +++ b/src/FL/Actors/Coordinator.hs @@ -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) diff --git a/src/FL/Actors/Evaluator.hs b/src/FL/Actors/Evaluator.hs new file mode 100644 index 0000000..f68357c --- /dev/null +++ b/src/FL/Actors/Evaluator.hs @@ -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) diff --git a/src/FL/Actors/Logger.hs b/src/FL/Actors/Logger.hs new file mode 100644 index 0000000..3babb7e --- /dev/null +++ b/src/FL/Actors/Logger.hs @@ -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 diff --git a/src/FL/Actors/PeerCoordinator.hs b/src/FL/Actors/PeerCoordinator.hs new file mode 100644 index 0000000..560db5a --- /dev/null +++ b/src/FL/Actors/PeerCoordinator.hs @@ -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) diff --git a/src/FL/Actors/Trainer.hs b/src/FL/Actors/Trainer.hs new file mode 100644 index 0000000..db637be --- /dev/null +++ b/src/FL/Actors/Trainer.hs @@ -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 diff --git a/src/FL/CRDT.hs b/src/FL/CRDT.hs new file mode 100644 index 0000000..8235971 --- /dev/null +++ b/src/FL/CRDT.hs @@ -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) diff --git a/src/FL/Data/PaySim.hs b/src/FL/Data/PaySim.hs new file mode 100644 index 0000000..f6fbe39 --- /dev/null +++ b/src/FL/Data/PaySim.hs @@ -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) diff --git a/src/FL/NN.hs b/src/FL/NN.hs new file mode 100644 index 0000000..68a84b9 --- /dev/null +++ b/src/FL/NN.hs @@ -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) diff --git a/src/FL/P2P.hs b/src/FL/P2P.hs new file mode 100644 index 0000000..05d016f --- /dev/null +++ b/src/FL/P2P.hs @@ -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 ===" diff --git a/src/FL/Provider.hs b/src/FL/Provider.hs new file mode 100644 index 0000000..3ada013 --- /dev/null +++ b/src/FL/Provider.hs @@ -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 diff --git a/src/FL/Retry.hs b/src/FL/Retry.hs new file mode 100644 index 0000000..e267c62 --- /dev/null +++ b/src/FL/Retry.hs @@ -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 diff --git a/src/FL/Types.hs b/src/FL/Types.hs new file mode 100644 index 0000000..17dd817 --- /dev/null +++ b/src/FL/Types.hs @@ -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