diff --git a/Benchmark.hs b/Benchmark.hs new file mode 100644 --- /dev/null +++ b/Benchmark.hs @@ -0,0 +1,66 @@ +import Control.Monad (replicateM_) +import qualified Network.Transport.TCP as TCP +import Control.Distributed.Process +import Control.Distributed.Process.Internal.Types (NodeId(..)) +import Control.Distributed.Process.Node +import Control.Concurrent (newEmptyMVar, takeMVar, putMVar) +import Control.Concurrent.Async +import Pipes +import qualified Pipes.Prelude as PP + +import TPar.Server +import TPar.Server.Types +import TPar.Types + +singleNode :: IO () +singleNode = do + Right transport <- TCP.createTransport "localhost" "0" TCP.defaultTCPParameters + node <- newLocalNode transport initRemoteTable + runProcess node $ do + iface <- runServer + replicateM_ 1 $ spawnLocal $ runRemoteWorker iface + replicateM_ 100 $ do + prod <- enqueueAndFollow iface jobReq + runEffect $ prod >-> PP.drain + +jobReq = JobRequest { jobName = JobName "hello" + , jobPriority = Priority 0 + , jobCommand = "echo" + , jobArgs = ["hello", "world"] + , jobCwd = "." + , jobEnv = Nothing + } + +multiNode :: IO () +multiNode = do + Right transport1 <- TCP.createTransport "localhost" "9000" TCP.defaultTCPParameters + node1 <- newLocalNode transport1 initRemoteTable + ifaceVar <- newEmptyMVar + doneVar <- newEmptyMVar + async $ runProcess node1 $ do + iface <- runServer + liftIO $ putMVar ifaceVar iface + liftIO $ takeMVar doneVar + iface <- takeMVar ifaceVar + + let nid = NodeId (TCP.encodeEndPointAddress "localhost" "9000" 0) + + Right transport2 <- TCP.createTransport "localhost" "0" TCP.defaultTCPParameters + node2 <- newLocalNode transport2 initRemoteTable + async $ runProcess node2 $ do + (sq, rq) <- newChan :: Process (SendPort ServerIface, ReceivePort ServerIface) + nsendRemote nid "tpar" sq + iface <- receiveChan rq + replicateM_ 1 $ spawnLocal $ runRemoteWorker iface + + Right transport3 <- TCP.createTransport "localhost" "0" TCP.defaultTCPParameters + node3 <- newLocalNode transport3 initRemoteTable + runProcess node3 $ do + replicateM_ 100 $ do + prod <- enqueueAndFollow iface jobReq + runEffect $ prod >-> PP.print + + putMVar doneVar () + +main :: IO () +main = multiNode diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright (c) 2013, Ben Gamari + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Ben Gamari nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/Main.hs b/Main.hs new file mode 100644 --- /dev/null +++ b/Main.hs @@ -0,0 +1,456 @@ +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE OverloadedStrings #-} + +import Control.Monad.Catch +import Control.Monad (when, unless, forever, replicateM_) +import Control.Monad.IO.Class +import Control.Error hiding (err) +import Data.Foldable +import qualified Data.Map.Strict as M +import Data.Time.Clock +import Data.Time.Format.Human +import System.Exit +import System.IO (stderr, stdout) +import Data.Aeson as Aeson +import qualified Data.ByteString.Lazy.Char8 as BS.L + +import qualified Text.PrettyPrint.ANSI.Leijen as T.PP +import Text.PrettyPrint.ANSI.Leijen (Doc, (<+>), (<$$>)) +import qualified Network.Transport.TCP as TCP +import Network.Socket (ServiceName, HostName) +import Network.BSD (getHostName) +import Options.Applicative hiding (action, header) +import qualified Options.Applicative as O +import Control.Concurrent (threadDelay) +import Control.Distributed.Process +import Control.Distributed.Process.Internal.Types (NodeId(..)) +import Control.Distributed.Process.Node +import qualified Text.Trifecta as TT + +import Pipes +import qualified Pipes.Concurrent as P.C +import qualified Pipes.Prelude as P.P + +import TPar.Rpc +import TPar.ProcessPipe hiding (runProcess) +import TPar.Server +import TPar.SubPubStream as SubPub +import TPar.Server.Types +import TPar.JobMatch +import TPar.Types + +portOption :: Mod OptionFields ServiceName -> Parser ServiceName +portOption m = + option str + ( short 'p' <> long "port" + <> value "5757" <> m + ) + +hostOption :: Parser HostName +hostOption = + strOption ( short 'H' <> long "host" <> value "localhost" <> help "server host name" ) + +jobMatchArg :: Parser JobMatch +jobMatchArg = argument (liftTrifecta parseJobMatch) (help "job match expression") + +type Mode = IO () + +tpar :: ParserInfo Mode +tpar = O.info (helper <*> tparParser) + $ O.fullDesc + <> O.progDesc "Start queues, add workers, and enqueue tasks" + <> O.header "tpar - simple distributed task queuing" + +tparParser :: Parser Mode +tparParser = + subparser + $ command "server" ( info modeServer + $ fullDesc <> progDesc "Start a server") + <> command "worker" ( info modeWorker + $ fullDesc <> progDesc "Start a worker") + <> command "enqueue" ( info modeEnqueue + $ fullDesc <> progDesc "Enqueue a job") + <> command "status" ( info modeStatus + $ fullDesc <> progDesc "Show queue status") + <> command "kill" ( info modeKill + $ fullDesc <> progDesc "Kill or dequeue a job") + <> command "rerun" ( info modeRerun + $ fullDesc <> progDesc "Restart a failed job") + <> command "watch" ( info modeWatch + $ fullDesc <> progDesc "Watch the output of a set of jobs") + <> command "dump" ( info modeDump + $ fullDesc <> progDesc "Dump details about a set of jobs in JSON") + +withServer' :: HostName -> ServiceName + -> (ServerIface -> Process a) -> Process a +withServer' host port action = do + let nid :: NodeId + nid = NodeId (TCP.encodeEndPointAddress host port 0) + -- request server interface + mref <- monitorNode nid + (sq, rq) <- newChan :: Process (SendPort ServerIface, ReceivePort ServerIface) + nsendRemote nid "tpar" sq + iface <- receiveWait + [ matchIf (\(NodeMonitorNotification ref _ _) -> ref == mref) $ + \(NodeMonitorNotification _ _ reason) -> + case reason of + DiedDisconnect -> fail "Failed to connect. Are you sure there is a server running?" + _ -> fail $ show reason + , matchChan rq pure + ] + + unless (protocolVersion iface == currentProtocolVersion) $ + fail $ "tpar protocol version mismatch: server speaks "++show (protocolVersion iface) + ++ ", client speaks "++show currentProtocolVersion + + -- request server interface + link (serverPid iface) + unlinkNode nid + r <- action iface + unlink (serverPid iface) + return r + +withServer :: HostName -> ServiceName + -> (ServerIface -> Process ()) -> IO () +withServer host port action = do + hostname <- getHostName + Right transport <- TCP.createTransport hostname "0" TCP.defaultTCPParameters + node <- newLocalNode transport initRemoteTable + runProcess node $ withServer' host port action + +modeWorker :: Parser Mode +modeWorker = + run <$> option (auto >>= checkNWorkers) + ( short 'N' <> long "workers" <> value 1 + <> help "number of local workers to start" + ) + <*> hostOption + <*> portOption (help "server port number") + <*> option (Just <$> (auto <|> pure 10)) + ( short 'r' <> long "reconnect" <> metavar "SECONDS" <> value Nothing + <> help "attempt to reconnect when server vanishes (with optional retry period); otherwise terminates on server vanishing" + ) + <* helper + where + checkNWorkers n + | n >= 1 = return n + | otherwise = fail "Worker count (-N) should be at least one" + + run nWorkers serverHost serverPort reconnectPeriod = + perhapsRepeat $ withServer serverHost serverPort $ \serverIface -> do + replicateM_ nWorkers $ spawnLocal $ runRemoteWorker serverIface + liftIO $ forever threadDelay maxBound + where + perhapsRepeat action + | Just period <- reconnectPeriod = forever $ do + handleAll (liftIO . print) action + liftIO (threadDelay $ 1000*1000*period) + | otherwise = action + +modeServer :: Parser Mode +modeServer = + run <$> option str (short 'H' <> long "host" <> value "localhost" + <> help "interface address to listen on" ) + <*> portOption (help "port to listen on") + <*> option auto ( short 'N' <> long "workers" <> value 0 + <> help "number of local workers to start" + ) + <* helper + where + run serverHost serverPort nLocalWorkers = do + Right transport <- TCP.createTransport serverHost serverPort TCP.defaultTCPParameters + node <- newLocalNode transport initRemoteTable + runProcess node $ do + iface <- runServer + replicateM_ nLocalWorkers $ spawnLocal $ runRemoteWorker iface + liftIO $ forever $ threadDelay maxBound + +modeEnqueue :: Parser Mode +modeEnqueue = + run <$> hostOption + <*> portOption (help "server port number") + <*> sinkType + <*> switch (short 'w' <> long "watch" <> help "Watch output of task") + <*> option (JobName <$> str) + (short 'n' <> long "name" <> value (JobName "unnamed-job") + <> help "Set the job's name") + <*> option str + (short 'd' <> long "directory" <> value "." + <> help "Set the directory the job will be launched from (relative to the working directory of the worker who runs it)") + <*> option (Priority <$> auto) + (short 'P' <> long "priority" <> value (Priority 0) + <> help "Set the job's priority") + <*> some (argument str idm) + <* helper + where + run :: HostName -> ServiceName + -> OutputStreams (Maybe FilePath) -> Bool + -> JobName -> FilePath -> Priority -> [String] + -> IO () + run serverHost serverPort sink watch name dir priority (cmd:args) = + withServer serverHost serverPort $ \iface -> do + let jobReq = JobRequest { jobName = name + , jobPriority = priority + , jobCommand = cmd + , jobArgs = args + , jobSinks = sink + , jobCwd = dir + , jobEnv = Nothing + } + if watch + then do + prod <- enqueueAndFollow iface jobReq + code <- runEffect $ prod >-> P.P.mapM_ (processOutputToHandles $ OutputStreams stdout stderr) + case code of + ExitSuccess -> return () + ExitFailure n -> + liftIO $ putStrLn $ "exited with code "++show n + else do + Right _ <- callRpc (enqueueJob iface) (jobReq, Nothing) + return () + + run _ _ _ _ _ _ _ _ = fail "Expected command line" + + sinkType :: Parser (OutputStreams (Maybe FilePath)) + sinkType = + OutputStreams + <$> option (Just <$> str) + (short 'o' <> long "output" <> metavar "FILE" + <> help "remote file to log standard output to" + <> value Nothing) + <*> option (Just <$> str) + (short 'e' <> long "error" <> metavar "FILE" + <> help "remote file to log standard error to" + <> value Nothing) + +liftTrifecta :: TT.Parser a -> ReadM a +liftTrifecta parser = do + s <- str + case TT.parseString parser mempty s of + TT.Success a -> return a + TT.Failure err -> fail $ show $ TT._errDoc err + +modeWatch :: Parser Mode +modeWatch = + run <$> hostOption + <*> portOption (help "server port number") + <*> jobMatchArg + <* helper + where + run serverHost serverPort jobMatch = + withServer serverHost serverPort $ \iface -> do + inputs <- subscribeToJobs iface jobMatch + let input = foldMap fold inputs + failed = M.keys $ M.filter isNothing inputs + unless (null failed) $ liftIO $ print + $ T.PP.red "warning: failed to attach to" + <+> prettyShow (length failed) <+> "jobs" + <$$> T.PP.nest 4 (T.PP.vcat $ map prettyShow failed) + runEffect $ P.C.fromInput input + >-> P.P.mapM_ (processOutputToHandles $ OutputStreams stdout stderr) + +subscribeToJobs :: ServerIface -> JobMatch + -> Process (M.Map JobId (Maybe (P.C.Input ProcessOutput))) +subscribeToJobs iface jobMatch = do + Right jobs <- callRpc (getQueueStatus iface) jobMatch + prods <- traverse SubPub.subscribe + $ M.unions + [ M.singleton jobId jobMonitor + | Job {..} <- jobs + , Running {..} <- pure jobState + ] + traverse (traverse producerToInput) prods + where + producerToInput :: Producer ProcessOutput Process ExitCode + -> Process (P.C.Input ProcessOutput) + producerToInput prod = do + (output, input) <- liftIO $ P.C.spawn (P.C.bounded 1) + void $ spawnLocal $ runEffect $ void prod >-> P.C.toOutput output + return input + +modeStatus :: Parser Mode +modeStatus = + run <$> hostOption + <*> portOption (help "server port number") + <*> switch (short 'v' <> long "verbose" <> help "verbose queue status") + <*> (jobMatchArg <|> pure (NegMatch NoMatch)) + <* helper + where + run serverHost serverPort verbose jobMatch = + withServer serverHost serverPort $ \iface -> do + Right jobs <- callRpc (getQueueStatus iface) jobMatch + time <- liftIO getCurrentTime + let prettyTime = T.PP.text . humanReadableTime' time + liftIO $ T.PP.putDoc $ T.PP.vcat $ map (prettyJob verbose prettyTime) jobs ++ [mempty] + +modeDump :: Parser Mode +modeDump = + run <$> hostOption + <*> portOption (help "server port number") + <*> (jobMatchArg <|> pure (NegMatch NoMatch)) + <* helper + where + run serverHost serverPort jobMatch = + withServer serverHost serverPort $ \iface -> do + Right jobs <- callRpc (getQueueStatus iface) jobMatch + liftIO $ BS.L.putStrLn $ Aeson.encode $ map jobToJson jobs + +jobToJson :: Job -> Value +jobToJson (Job {jobRequest = JobRequest{..}, ..}) = + object [ "id" .= fromEnum jobId + , "name" .= case jobName of JobName s -> s + , "priority" .= fromEnum jobPriority + , "command" .= jobCommand + , "arguments" .= jobArgs + , "working_dir" .= jobCwd + , "remote_logs" .= object [ "stdout" .= stdOut jobSinks + , "stderr" .= stdErr jobSinks + ] + , "environment" .= jobEnv + , "state" .= state + ] + where + stateObj s rest = object $ ("state" .= (s :: String)) : rest + state = + case jobState of + Queued {..} -> + stateObj "queued" + [ "queued_at" .= jobQueueTime ] + Starting {..} -> + stateObj "starting" + [ "queued_at" .= jobQueueTime + , "starting_since" .= jobStartingTime + , "running_on" .= show jobProcessId + ] + Running {..} -> + stateObj "running" + [ "queued_at" .= jobQueueTime + , "started_at" .= jobStartTime + , "running_on" .= show jobProcessId + ] + Finished {..} -> + stateObj "finished" + [ "queued_at" .= jobQueueTime + , "started_at" .= jobStartTime + , "finished_at" .= jobFinishTime + , "ran_on" .= show jobWorkerNode + , "exit_code" .= getExitCode jobExitCode + ] + Failed {..} -> + stateObj "finished" + [ "queued_at" .= jobQueueTime + , "started_at" .= jobStartTime + , "failed_at" .= jobFailedTime + , "ran_on" .= show jobWorkerNode + , "error" .= jobErrorMsg + ] + Killed {..} -> + stateObj "killed" + [ "queued_at" .= jobQueueTime + , "started_at" .= jobKilledStartTime + , "killed_at" .= jobKilledTime + , "ran_on" .= show jobKilledWorkerNode + ] + +prettyJob :: Bool -> (UTCTime -> Doc) -> Job -> Doc +prettyJob verbose prettyTime (Job {..}) = + T.PP.vcat $ [header] ++ (if verbose then [details] else []) + where + JobRequest {..} = jobRequest + + twoCols :: Int -> [(Doc, Doc)] -> Doc + twoCols width = + T.PP.vcat . map (\(a,b) -> T.PP.fillBreak width a <+> T.PP.align b) + + header = + T.PP.fillBreak 5 (prettyJobId jobId) + <+> T.PP.fillBreak 50 (prettyJobName jobName) + <+> prettyJobState jobState + + details = + T.PP.indent 4 $ twoCols 15 + [ ("priority:", prettyPriority jobPriority) + , ("queued:", prettyTime (jobQueueTime jobState)) + , ("command:", T.PP.text jobCommand) + , ("arguments:", T.PP.hsep $ map T.PP.text jobArgs) + , ("logging:", T.PP.vcat $ toList $ (\x y -> x<>":"<+>y) + <$> OutputStreams "stdout" "stderr" + <*> fmap (maybe "none" T.PP.text) jobSinks) + , ("status:", prettyDetailedState jobState) + ] + <$$> mempty + + prettyDetailedState Queued{..} = + "waiting to run" <+> T.PP.parens ("since" <+> prettyTime jobQueueTime) + prettyDetailedState Starting{..} = + "starting on" <+> prettyShow jobProcessId + <+> T.PP.parens ("since" <+> prettyTime jobStartingTime) + prettyDetailedState Running{..} = + "running on" <+> prettyShow jobProcessId + <+> T.PP.parens ("since" <+> prettyTime jobStartTime) + prettyDetailedState Finished{..} = + "finished with" <+> prettyShow (getExitCode jobExitCode) + <+> T.PP.parens ("at" <+> prettyTime jobFinishTime) + <$$> "started at" <+> prettyTime jobStartTime + <$$> "ran on" <+> prettyShow jobWorkerNode + prettyDetailedState Failed{..} = + "failed with error" <+> T.PP.parens (prettyTime jobFailedTime) + <$$> "started at" <+> prettyTime jobStartTime + <$$> T.PP.indent 4 (T.PP.text jobErrorMsg) + prettyDetailedState Killed{..} = + "killed at user request" <+> T.PP.parens (prettyTime jobKilledTime) + <$$> maybe "never started" (\t -> "started at" <+> prettyTime t) jobKilledStartTime + + prettyJobState Queued{} = T.PP.blue "queued" + prettyJobState Starting{} = T.PP.blue "starting" + prettyJobState Running{} = T.PP.cyan "running" + prettyJobState Finished{..} + | ExitFailure _ <- jobExitCode = T.PP.yellow $ "failed" + | otherwise = T.PP.green "finished" + prettyJobState Failed{} = T.PP.red "job failed" + prettyJobState Killed{} = T.PP.yellow "killed" + + prettyJobId (JobId n) = T.PP.int n + prettyJobName (JobName name) = T.PP.bold $ T.PP.text name + prettyPriority (Priority p) = T.PP.int p + +getExitCode :: ExitCode -> Int +getExitCode ExitSuccess = 0 +getExitCode (ExitFailure c) = c + +prettyShow :: Show a => a -> Doc +prettyShow = T.PP.text . show + +modeKill :: Parser Mode +modeKill = + run <$> hostOption + <*> portOption (help "server port number") + <*> jobMatchArg + <* helper + where + run serverHost serverPort jobMatch = + withServer serverHost serverPort $ \iface -> do + Right jobs <- callRpc (killJobs iface) jobMatch + liftIO $ T.PP.putDoc $ T.PP.vcat $ map (prettyJob False prettyShow) jobs ++ [mempty] + liftIO $ when (null jobs) $ exitWith $ ExitFailure 1 + +modeRerun :: Parser Mode +modeRerun = + run <$> hostOption + <*> portOption (help "server port number") + <*> jobMatchArg + <* helper + where + run serverHost serverPort jobMatch = + withServer serverHost serverPort $ \iface -> do + Right jobs <- callRpc (rerunJobs iface) jobMatch + liftIO $ T.PP.putDoc $ T.PP.vcat $ map (prettyJob False prettyShow) jobs ++ [mempty] + liftIO $ when (null jobs) $ exitWith $ ExitFailure 1 + +main :: IO () +main = do + run <- execParser tpar + res <- runExceptT $ tryIO run + case res of + Left err -> putStrLn $ "error: "++show err + Right () -> return () diff --git a/README.mkd b/README.mkd new file mode 100644 --- /dev/null +++ b/README.mkd @@ -0,0 +1,75 @@ +# tpar — simple parallel job scheduling + +`tpar` is a simple tool for concurrent job scheduling. Say you have a +directory full of files which need processing, + +```bash +$ ls +file1 file2 file3 file4 file5 +... +``` + +Usually one could use a `bash` for loop, + +```bash +$ for f in *; do process $f; done; +``` + +But if `process` is a long-running task and you have many cores at +your disposal, it would be nice to speed things up a bit, + +```bash +$ tpar server -N8 +$ for f in *; do tpar enqueue -- process $f; done; +``` + +If you have multiple machines with the data mounted over, say, NFS, they can +also help with churning through the queue, + +```bash +$ for m in worker1 worker2 worker3; do +> ssh $m -- tpar worker -H`hostname`; +> done +``` + +## Commands + +`tpar` has several subcommands, + + * `tpar server` starts a local queue server. + * `tpar worker` starts a worker associated with + the given queue + * `tpar enqueue -- $cmd` enqueues a job in the given queue + * `tpar status` allows you to query for the status of the queue. You can also provide a job match expression + * `tpar kill` kills a running task (specified by a job match expression) + * `tpar watch` is analogous `tail -f`, watching the output of a set of running tasks + * `tpar dump` dumps a JSON representation of the queue state. + +Nearly all of these commands will require that the `-H` option be provided +specifying the canonical hostname of the queue server (the machine running +`tpar server`). + + +## Job match expressions + +Several `tpar` commands accept a *job match expression*, which specifies the +subset of jobs on which the command should act. For instance (note the quotes to +ensure that `bash` doesn't interpret our symbols), + +``` +$ tpar status '*' # This is equivalent to `tpar status` run without an argument +$ tpar status id=2 +$ tpar status state=running +$ tpar status 'name="my-job" or name="my-other-job"' +``` + +These expressions consist of the primitive matches, + * `name="STRING"`, which matches on the job name provided in the `--name` of + `tpar enqueue`. + * `id=`, which matches on the job ID + * `state=`, which matches on the current state of the job (`queued`, `running`, + `finished`, `failed`, `killed`, or `code=N`) + * `*`, which matches all jobs + +These matches can be connected with the `and` and `or` operators, and inverted +with `!`. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/TPar/JobMatch.hs b/TPar/JobMatch.hs new file mode 100644 --- /dev/null +++ b/TPar/JobMatch.hs @@ -0,0 +1,144 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE NamedFieldPuns #-} + +module TPar.JobMatch where + +import Data.Foldable (traverse_) +import Control.Monad (void) +import Control.Applicative ((<|>)) +import System.Exit +import GHC.Generics (Generic) + +import Data.Binary +import Text.Trifecta +import Text.Parser.Expression +import Text.Parser.Token.Style (emptyOps) + +import TPar.Types + +data GlobAtom = WildCard + | Literal String + deriving (Generic, Show) + +instance Binary GlobAtom + +type Glob = [GlobAtom] + +globAtomToParser :: GlobAtom -> Parser () +globAtomToParser WildCard = void anyChar +globAtomToParser (Literal a) = void $ string a + +globToParser :: Glob -> Parser () +globToParser atoms = traverse_ globAtomToParser atoms >> eof + +globMatches :: Glob -> String -> Bool +globMatches glob str = + case parseString (globToParser glob) mempty str of + Success () -> True + Failure _ -> False + +parseGlob :: Parser Glob +parseGlob = many $ wildCard <|> literal + where + wildCard = char '*' >> pure WildCard + literal = Literal <$> some (noneOf reserved) + + reserved = "*\"" + +data JobMatch = NoMatch + | NegMatch JobMatch + | NameMatch Glob + | JobIdMatch JobId + | AltMatch [JobMatch] + | AllMatch [JobMatch] + | StateMatch StateMatch + deriving (Generic, Show) + +instance Binary JobMatch + +data StateMatch = IsQueued + | IsRunning + | IsFinished (Maybe ExitCode) + | IsFailed + | IsKilled + deriving (Generic, Show) + +instance Binary StateMatch + +jobMatches :: JobMatch -> Job -> Bool +jobMatches NoMatch _ = False +jobMatches (NegMatch m) job = not $ jobMatches m job +jobMatches (NameMatch glob) job = globMatches glob name + where JobName name = jobName $ jobRequest job +jobMatches (JobIdMatch jobid) job = jobId job == jobid +jobMatches (AltMatch alts) job = any (`jobMatches` job) alts +jobMatches (AllMatch alts) job = all (`jobMatches` job) alts +jobMatches (StateMatch s) job = stateMatches s (jobState job) + +stateMatches :: StateMatch -> JobState -> Bool +stateMatches IsQueued Queued{} = True +stateMatches IsRunning Running{} = True +stateMatches (IsFinished Nothing) Finished{} = True +stateMatches (IsFinished (Just code)) Finished {jobExitCode} = code == jobExitCode +stateMatches IsFailed Failed{} = True +stateMatches IsKilled Killed{} = True +stateMatches _ _ = False + +opTable :: OperatorTable Parser JobMatch +opTable = + [ [ prefix "!" NegMatch ] + , [ binary "&" (\x y -> AllMatch [x,y]) AssocLeft + , binary "and" (\x y -> AllMatch [x,y]) AssocLeft + ] + , [ binary "|" (\x y -> AltMatch [x,y]) AssocLeft + , binary "or" (\x y -> AltMatch [x,y]) AssocLeft + ] + ] + where + binary name fun assoc = Infix (fun <$ reservedOp name) assoc + prefix name fun = Prefix (fun <$ reservedOp name) + reservedOp name = reserve emptyOps name + +parseJobMatch :: Parser JobMatch +parseJobMatch = expr + where + expr = buildExpressionParser opTable term "job match expression" + + term = parens expr <|> simple + + simple = starMatch <|> jobIdMatch <|> stateMatch <|> nameMatch + + starMatch :: Parser JobMatch + starMatch = char '*' >> pure (NegMatch NoMatch) + + nameMatch :: Parser JobMatch + nameMatch = do + void $ string "name=" + NameMatch <$> between (char '"') (char '"') parseGlob + + stateMatch :: Parser JobMatch + stateMatch = do + void $ string "state=" + StateMatch <$> parseJobState + + jobIdMatch :: Parser JobMatch + jobIdMatch = do + void $ string "id=" + JobIdMatch . JobId . fromIntegral <$> integer + +parseJobState :: Parser StateMatch +parseJobState = + (string "queued" *> pure IsQueued) + <|> (string "running" *> pure IsRunning) + <|> (string "finished" *> pure (IsFinished Nothing)) + <|> (string "done" *> pure (IsFinished Nothing)) + <|> (string "succeeded" *> pure (IsFinished (Just ExitSuccess))) + <|> (string "code=" *> (IsFinished . Just <$> exitCode)) + <|> (string "failed" *> pure IsFailed) + <|> (string "killed" *> pure IsKilled) + where + exitCode = ( "exit code") $ do + n <- integer + pure $ case n of + 0 -> ExitSuccess + _ -> ExitFailure (fromIntegral n) diff --git a/TPar/ProcessPipe.hs b/TPar/ProcessPipe.hs new file mode 100644 --- /dev/null +++ b/TPar/ProcessPipe.hs @@ -0,0 +1,140 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module TPar.ProcessPipe ( ProcessOutput(..) + , runProcess + -- * Killing the process + , ProcessKilled(..) + -- * Deinterleaving output + , processOutputToHandles + , selectStream + , OutputStreams(..) + ) where + +import Control.Applicative +import Data.Monoid +import Data.Traversable +import qualified Pipes.Prelude as PP +import qualified Pipes.ByteString as PBS +import Data.ByteString (ByteString) +import qualified Data.ByteString as BS +import Control.Monad (msum) +import Control.Exception (Exception) +import System.IO (Handle) +import System.Exit + +import Pipes +import Pipes.Safe () -- for MonadCatch instance +import qualified Pipes.Concurrent as PC +import System.Process (runInteractiveProcess, ProcessHandle, waitForProcess, terminateProcess) +import Control.Concurrent.STM +import Control.Distributed.Process +import Control.Monad.Catch (handle, throwM) + +import Data.Binary +import Data.Binary.Put +import Data.Binary.Get +import GHC.Generics +import TPar.Utils + +processPipes :: MonadIO m + => FilePath -- ^ Executable name + -> [String] -- ^ Arguments + -> Maybe FilePath -- ^ Working directory + -> Maybe [(String,String)] -- ^ Optional environment + -> IO ( Consumer ByteString m () + , Producer ByteString m () + , Producer ByteString m () + , ProcessHandle) +processPipes cmd args cwd env = do + (stdin, stdout, stderr, phandle) <- runInteractiveProcess cmd args cwd env + return (PBS.toHandle stdin, PBS.fromHandle stdout, PBS.fromHandle stderr, phandle) + +data InterleaverCanTerminate = InterleaverCanTerminate deriving (Generic) +instance Binary InterleaverCanTerminate + +data InterleaveException = InterleaveException String + deriving (Show) +instance Exception InterleaveException + +interleave :: forall a. [Producer a Process ()] -> Producer a Process () +interleave producers = do + inputs <- lift $ forM producers $ \prod -> do + (output, input, seal) <- liftIO $ PC.spawn' (PC.bounded 10) + pid <- spawnLocal $ runEffect $ do + prod >-> PC.toOutput output + liftIO $ atomically seal + + _ <- monitor pid + return input + + let matchTermination = match $ \(ProcessMonitorNotification _ _pid reason) -> + case reason of + DiedNormal -> return Nothing + _ -> throwM $ InterleaveException $ show reason + matchData = matchSTM (PC.recv $ msum inputs) pure + go :: Producer a Process () + go = do + mx <- lift $ receiveWait [ matchTermination, matchData ] + case mx of + Nothing -> return () + Just x -> yield x >> go + go + +data ProcessOutput + = PutStdout !ByteString + | PutStderr !ByteString + deriving (Show, Generic) + +instance Binary ProcessOutput + +data OutputStreams a = OutputStreams { stdOut, stdErr :: a } + deriving (Show, Functor, Generic) +instance Binary a => Binary (OutputStreams a) +instance Foldable OutputStreams where + foldMap f (OutputStreams x y) = f x <> f y +instance Applicative OutputStreams where + pure x = OutputStreams x x + OutputStreams f g <*> OutputStreams x y = OutputStreams (f x) (g y) + +-- Unfortunate orphan +instance Binary ExitCode where + get = do + code <- getInt32le + return $ case code of + 0 -> ExitSuccess + _ -> ExitFailure (fromIntegral code) + put ExitSuccess = putInt32le 0 + put (ExitFailure code) = putInt32le (fromIntegral code) + +data ProcessKilled = ProcessKilled + deriving (Show, Generic) + +instance Binary ProcessKilled +instance Exception ProcessKilled + +selectStream :: OutputStreams (ByteString -> a) -> ProcessOutput -> a +selectStream (OutputStreams out err) outs = + case outs of + PutStdout bs -> out bs + PutStderr bs -> err bs + +runProcess :: FilePath -> [String] -> Maybe FilePath + -> Producer ProcessOutput Process ExitCode +runProcess cmd args cwd = do + lift $ tparDebug "starting process" + (_stdin, stdout, stderr, phandle) <- liftIO $ processPipes cmd args cwd Nothing + let processKilled ProcessKilled = liftIO $ do + terminateProcess phandle + throwM ProcessKilled + handle processKilled $ do + interleave [ stderr >-> PP.map PutStderr + , stdout >-> PP.map PutStdout + ] + liftIO $ waitForProcess phandle + +processOutputToHandles :: MonadIO m + => OutputStreams Handle -> ProcessOutput -> m () +processOutputToHandles handles = + selectStream $ fmap (\hdl bs -> liftIO $ BS.hPut hdl bs) handles diff --git a/TPar/Rpc.hs b/TPar/Rpc.hs new file mode 100644 --- /dev/null +++ b/TPar/Rpc.hs @@ -0,0 +1,50 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} + +module TPar.Rpc where + +import Control.Distributed.Process +import Control.Distributed.Process.Serializable +import Data.Binary + +newtype RpcSendPort a b = RpcSendPort (SendPort (a, SendPort b)) + deriving (Show, Binary, Serializable) +newtype RpcRecvPort a b = RpcRecvPort (ReceivePort (a, SendPort b) ) + +newRpc :: (Serializable a, Serializable b) + => Process (RpcSendPort a b, RpcRecvPort a b) +newRpc = do + (sp, rp) <- newChan + return (RpcSendPort sp, RpcRecvPort rp) + +-- | Call a remote procedure or return 'Left' if the handling process is no +-- longer alive. +callRpc :: (Serializable a, Serializable b) + => RpcSendPort a b -> a -> Process (Either DiedReason b) +callRpc (RpcSendPort sp) x = do + (reply_sp, reply_rp) <- newChan + mref <- monitorPort sp + sendChan sp (x, reply_sp) + receiveWait [ matchIf (\(PortMonitorNotification mref' _ _) -> mref == mref') + (\(PortMonitorNotification _ _ reason) -> pure $ Left reason) + , matchChan reply_rp (pure . Right) + ] + +matchRpc :: (Serializable a, Serializable b) + => RpcRecvPort a b -> (a -> Process (b, c)) -> Match c +matchRpc rp handler = matchRpc' rp $ \_ x reply -> do + (y, z) <- handler x + reply y + return z + +type RpcHandler a b c + = ProcessId -- ^ 'ProcessId' of the requestor + -> a -- ^ arguments + -> (b -> Process ()) -- ^ reply action + -> Process c -- ^ return + +-- | Allow deferred replies +matchRpc' :: (Serializable a, Serializable b) + => RpcRecvPort a b -> RpcHandler a b c -> Match c +matchRpc' (RpcRecvPort rp) handler = matchChan rp $ \(x, reply_sp) -> do + let requestor_pid = sendPortProcessId $ sendPortId reply_sp + handler requestor_pid x (sendChan reply_sp) diff --git a/TPar/Server.hs b/TPar/Server.hs new file mode 100644 --- /dev/null +++ b/TPar/Server.hs @@ -0,0 +1,372 @@ +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE NamedFieldPuns #-} + +module TPar.Server + ( -- * Workers + Worker + , localWorker + , sshWorker + , runRemoteWorker + -- * Running the server + , server + , runServer + -- * Convenience wrappers + , enqueueAndFollow + ) where + +import Control.Error +import Control.Applicative +import Control.Monad (void, forever, filterM) +import Data.Foldable +import Data.Traversable +import qualified Data.Heap as H +import qualified Data.Map as M +import qualified Data.Set as S +import qualified Data.ByteString as BS +import Control.Monad.Catch +import Control.Distributed.Process hiding (finally, bracket) +import Data.Time.Clock + +import System.IO ( openFile, hClose, IOMode(..), Handle ) +import System.Exit +import Control.Concurrent.STM + +import Pipes +import qualified Pipes.Prelude as P.P + +import TPar.Rpc +import TPar.SubPubStream as SubPub +import TPar.ProcessPipe +import TPar.Server.Types +import TPar.Types +import TPar.JobMatch +import TPar.Utils + +----------------------------------------------- +-- Convenience wrappers + +enqueueAndFollow :: ServerIface -> JobRequest + -> Process (Producer ProcessOutput Process ExitCode) +enqueueAndFollow iface jobReq = do + (rpcSp, rpcRp) <- newRpc + _jobId <- callRpc (enqueueJob iface) (jobReq, Just rpcSp) + receiveWait + [ matchRpc rpcRp $ \subPub -> do + mprod <- subscribe subPub + case mprod of + Just prod -> return ((), prod) + Nothing -> fail "enqueueAndFollow: Initial subscription failed" + ] + +----------------------------------------------- +-- Workers + +type Worker = JobRequest -> Producer ProcessOutput Process ExitCode + +localWorker :: Worker +localWorker req = runProcess (jobCommand req) (jobArgs req) (Just $ jobCwd req) + +sshWorker :: String -> FilePath -> Worker +sshWorker host rootPath req = do + runProcess "ssh" ([host, "--", "cd", cwd, ";", jobCommand req]++jobArgs req) Nothing + where + cwd = rootPath ++ "/" ++ jobCwd req -- HACK + +runRemoteWorker :: ServerIface -> Process () +runRemoteWorker (ServerIface {..}) = forever runOneJob + where + runOneJob = do + doneVar <- liftIO newEmptyTMVarIO + -- We run each process in a separate thread to ensure that ProcessKilled + -- exceptions go to the Process that is actually running the job being + -- killed + let finished = liftIO $ atomically $ putTMVar doneVar () + _pid <- spawnLocal $ flip finally finished $ do + Right (job, notifyJobStarted, jobFinishedSp) <- callRpc requestJob () + tparDebug $ "have job "++show (jobId job) + result <- runJobWithWorker job notifyJobStarted localWorker + finishedTime <- liftIO getCurrentTime + sendChan jobFinishedSp (finishedTime, result) + + liftIO $ atomically $ takeTMVar doneVar + +runJobWithWorker :: Job + -> JobStartedNotify + -> Worker + -> Process ExitCode +runJobWithWorker (Job {..}) notifyJobStarted worker = + withOutFiles (jobSinks jobRequest) $ \intoSinks -> do + (start, subPub, readRet) <- + SubPub.fromProducer' $ worker jobRequest >-> intoSinks + -- wait for confirmation from server before starting + Right () <- callRpc notifyJobStarted subPub + start + either throwM pure =<< liftIO (atomically readRet) + where + withOutFiles :: (MonadMask m, MonadIO m) + => OutputStreams (Maybe FilePath) + -> (Pipe ProcessOutput ProcessOutput m r -> m a) + -> m a + withOutFiles (OutputStreams (Just outPath) (Just errPath)) action + | outPath == errPath + = withOutFile outPath $ \hOut -> + action $ P.P.chain + $ processOutputToHandles (OutputStreams hOut hOut) + + | otherwise + = withOutFile outPath $ \hOut -> + withOutFile errPath $ \hErr -> + action $ P.P.chain $ processOutputToHandles (OutputStreams hOut hErr) + + withOutFiles (OutputStreams (Just outPath) Nothing) action + = withOutFile outPath $ \hOut -> + action $ P.P.chain $ toHandles $ OutputStreams (Just hOut) Nothing + + withOutFiles (OutputStreams Nothing (Just errPath)) action + = withOutFile errPath $ \hErr -> + action $ P.P.chain $ toHandles $ OutputStreams Nothing (Just hErr) + + withOutFiles _ action = action cat + + toHandles :: MonadIO m + => OutputStreams (Maybe Handle) + -> ProcessOutput -> m () + toHandles streams = liftIO . selectStream writeIt + where + writeIt :: OutputStreams (BS.ByteString -> IO ()) + writeIt = fmap (maybe (const $ return ()) BS.hPutStr) streams + + withOutFile :: (MonadIO m, MonadMask m) + => FilePath + -> (Handle -> m a) + -> m a + withOutFile path = bracket (liftIO $ openFile path WriteMode) (liftIO . hClose) + +------------------------------------------------ +-- the server + +-- | Spawn a process running a server +runServer :: Process ServerIface +runServer = do + q <- liftIO newJobQueue + iface <- server q + announce <- spawnLocal $ forever $ do + x <- expect :: Process (SendPort ServerIface) + sendChan x iface + register "tpar" announce + return iface + +-- | The heart of the server +server :: JobQueue -> Process ServerIface +server jobQueue = do + let protocolVersion = currentProtocolVersion + (enqueueJob, enqueueJobRp) <- newRpc + (requestJob, requestJobRp) <- newRpc + (getQueueStatus, getQueueStatusRp) <- newRpc + (killJobs, killJobsRp) <- newRpc + (rerunJobs, rerunJobsRp) <- newRpc + + serverPid <- spawnLocal $ void $ forever $ do + serverPid <- getSelfPid + receiveWait + [ matchRpc enqueueJobRp $ \(jobReq, notifyOriginator) -> do + tparDebug "enqueue" + jobQueueTime <- liftIO getCurrentTime + liftIO $ atomically $ do + jobId <- getFreshJobId jobQueue + queueJob jobQueue jobId jobQueueTime jobReq notifyOriginator + return (jobId, ()) + + , matchRpc' requestJobRp $ \workerPid () reply -> do + tparDebug $ "request job: "++show workerPid + _ <- spawnLocal $ do + link serverPid + handleJobRequest jobQueue workerPid reply + return () + + , matchRpc getQueueStatusRp $ \jobMatch -> do + q <- liftIO $ atomically $ getJobs jobQueue + let filtered = filter (jobMatches jobMatch) q + return (filtered, ()) + + , matchRpc killJobsRp $ \jobMatch -> do + killedJobs <- handleKillJobs jobQueue jobMatch + return (killedJobs, ()) + + , matchRpc rerunJobsRp $ \jobMatch -> do + reran <- handleRerunJobs jobQueue jobMatch + return (reran, ()) + ] + return $ ServerIface {..} + +handleJobRequest :: JobQueue + -> ProcessId + -> ((Job, JobStartedNotify, JobFinishedChan) -> Process ()) + -> Process () +handleJobRequest jobQueue workerPid reply = do + -- get a job + monRef <- monitor workerPid + job <- liftIO $ atomically $ takeQueuedJob jobQueue + + -- send the job to worker + (startedSp, startedRp) <- newRpc + (finishedSp, finishedRp) <- newChan + reply (job, startedSp, finishedSp) + let jobid = jobId job + jobProcessId = workerPid + jobWorkerNode = processNodeId workerPid + Queued {jobQueueTime} = jobState job + jobStartingTime <- liftIO getCurrentTime + liftIO $ atomically $ setJobState jobQueue jobid (Starting {..}) + tparDebug $ "job "++show jobid++" starting" + + -- wait for started notification + jobMonitor <- receiveWait + [matchRpc startedRp $ \jobMonitor -> do + forM_ (jobStartingNotify job) $ \notifyOriginator -> do + res <- callRpc notifyOriginator jobMonitor + case res of + Right () -> return () + Left reason -> + say $ "Job starting notification to "++show notifyOriginator + ++" failed due to "++show reason++". Starting anyways." + return ((), jobMonitor) + ] + jobStartTime <- liftIO getCurrentTime + liftIO $ atomically $ setJobState jobQueue jobid (Running {..}) + tparDebug $ "job "++show jobid++" running" + + -- wait for result + receiveWait + [ matchChan finishedRp $ \(jobFinishTime, jobExitCode) -> + liftIO $ atomically $ setJobState jobQueue jobid (Finished {..}) + + , matchIf (\(ProcessMonitorNotification ref _ _) -> ref == monRef) $ + \(ProcessMonitorNotification _ _ reason) -> do + tparDebug $ "job "++show jobid++" failed" + jobFailedTime <- liftIO getCurrentTime + let jobErrorMsg = show reason + liftIO $ atomically $ setJobState jobQueue jobid (Failed {..}) + ] + unmonitor monRef + +handleKillJobs :: JobQueue -> JobMatch -> Process [Job] +handleKillJobs jq@(JobQueue {..}) jobMatch = do + let shouldBeKilled :: Job -> Maybe (Maybe ProcessId, JobId) + shouldBeKilled job@(Job {..}) + | Running {..} <- jobState + , jobMatches jobMatch job = Just (Just jobProcessId, jobId) + | Queued {..} <- jobState + , jobMatches jobMatch job = Just (Nothing, jobId) + | otherwise = Nothing + jobsToKill <- liftIO $ atomically $ mapMaybe shouldBeKilled <$> getJobs jq + say $ "killing "++show jobsToKill + jobKilledTime <- liftIO getCurrentTime + killed <- forM jobsToKill $ \(pid, jobid) -> do + maybe (return ()) (flip exit ProcessKilled) pid + liftIO $ atomically $ do + oldState <- updateJob jq jobid $ \job -> + let state' = case jobState job of + Queued {..} -> Killed { jobKilledStartTime = Nothing + , jobKilledWorkerNode = Nothing + , .. + } + Running {..} -> Killed { jobKilledStartTime = Just jobStartTime + , jobKilledWorkerNode = Just $ processNodeId jobProcessId + , .. + } + s -> s + in (job {jobState=state'}, jobState job) + case oldState of + Queued {} -> do + modifyTVar jobQueue $ H.fromList . filter (\(_, job') -> job' /= jobid) . toList + return $ Just jobid + Running {} -> return $ Just jobid + _ -> return Nothing + let killedSet = S.fromList $ catMaybes killed + liftIO $ atomically $ filter (\job -> jobId job `S.member` killedSet) <$> getJobs jq + +handleRerunJobs :: JobQueue -> JobMatch -> Process [Job] +handleRerunJobs jq@(JobQueue {..}) jobMatch = do + jobQueueTime <- liftIO getCurrentTime + liftIO $ atomically $ do + allJobs <- getJobs jq + filterM (\job -> reEnqueueJob jq job jobQueueTime) + $ filter (jobMatches jobMatch) allJobs + +----------------------------------------------------- +-- primitives + +-- | Our job queue state +data JobQueue = JobQueue { freshJobIds :: TVar [JobId] + , jobQueue :: TVar (H.Heap (Priority, JobId)) + , jobs :: TVar (M.Map JobId Job) + } + +getFreshJobId :: JobQueue -> STM JobId +getFreshJobId (JobQueue {..}) = do + x:xs <- readTVar freshJobIds + writeTVar freshJobIds xs + return x + +newJobQueue :: IO JobQueue +newJobQueue = JobQueue <$> newTVarIO [ JobId i | i <- [0..] ] + <*> newTVarIO H.empty + <*> newTVarIO mempty + +takeQueuedJob :: JobQueue -> STM Job +takeQueuedJob jq@(JobQueue {..}) = do + q <- readTVar jobQueue + case H.viewMin q of + Nothing -> retry + Just ((_, jobid), q') -> do + writeTVar jobQueue q' + getJob jq jobid + +getJob :: JobQueue -> JobId -> STM Job +getJob (JobQueue {..}) jobId = (M.! jobId) <$> readTVar jobs + +updateJob :: JobQueue -> JobId -> (Job -> (Job, a)) -> STM a +updateJob (JobQueue {..}) jobId f = do + jobsMap <- readTVar jobs + Just x <- pure $ M.lookup jobId jobsMap + let (x', r) = f x + writeTVar jobs $ M.insert jobId x' jobsMap + return r + +setJobState :: JobQueue -> JobId -> JobState -> STM () +setJobState jobQueue jobId newState = + updateJob jobQueue jobId (\s -> (s {jobState = newState}, ())) + +-- | Place a finished, failed, or killed job back on the run queue. +-- Returns 'True' if re-queued. +reEnqueueJob :: JobQueue -> Job -> UTCTime -> STM Bool +reEnqueueJob jq Job{..} reEnqueueTime + | reQueuable = do + setJobState jq jobId (Queued reEnqueueTime) + modifyTVar (jobQueue jq) $ H.insert (jobPriority jobRequest, jobId) + return True + | otherwise = return False + where + reQueuable = + case jobState of + Queued {} -> False + Starting {} -> False + Running {} -> False + Failed {} -> True + Finished {} -> True + Killed {} -> True + +queueJob :: JobQueue -> JobId -> UTCTime -> JobRequest + -> Maybe JobStartingNotify + -> STM () +queueJob (JobQueue {..}) jobId jobQueueTime jobRequest jobStartingNotify = do + modifyTVar jobQueue $ H.insert (prio, jobId) + modifyTVar jobs $ M.insert jobId (Job {jobState = Queued {..}, ..}) + where + prio = jobPriority jobRequest + +getJobs :: JobQueue -> STM [Job] +getJobs (JobQueue {..}) = + M.elems <$> readTVar jobs diff --git a/TPar/Server/Types.hs b/TPar/Server/Types.hs new file mode 100644 --- /dev/null +++ b/TPar/Server/Types.hs @@ -0,0 +1,78 @@ +{-# LANGUAGE DeriveGeneric #-} + +module TPar.Server.Types where + +import Control.Distributed.Process +import Data.Binary +import Data.Time.Clock +import System.Exit +import GHC.Generics + +import TPar.Rpc +import TPar.JobMatch +import TPar.Types +import TPar.SubPubStream +import TPar.ProcessPipe + +-- | A channel which should be used by a 'Worker' to indicate that a job has +-- been started. It provides a 'SendPort' which will be called by the server to indicate that any atomic watches have been established and that the job can proceed. +type JobStartedNotify = RpcSendPort (SubPubSource ProcessOutput ExitCode) () + +-- | A channel which should be used by a 'Worker' to indicate that a job has +-- finished. +type JobFinishedChan = SendPort (UTCTime, ExitCode) + +data ServerIface = + ServerIface { protocolVersion :: ProtocolVersion + , serverPid :: ProcessId + , enqueueJob :: RpcSendPort (JobRequest, Maybe JobStartingNotify) JobId + , requestJob :: RpcSendPort () (Job, JobStartedNotify, JobFinishedChan) + , killJobs :: RpcSendPort JobMatch [Job] + , getQueueStatus :: RpcSendPort JobMatch [Job] + , rerunJobs :: RpcSendPort JobMatch [Job] + } + deriving (Generic) + +instance Binary ServerIface + +type ProtocolVersion = Int + +currentProtocolVersion :: ProtocolVersion +currentProtocolVersion = 5 + +{- $ the-story + +Starting a job (in the general case with an atomic monitor) happens as follows, + +1. An 'enqueueJob' request is sent to the server. This request carries a +'JobStartingNotify' which will be used to guarantee that the job doesn't start +before the monitor is established. + +2. The server places the job on the queue in the 'Queued' state. + +3. A worker sends a 'requestJob' request to the server, the job is de-queued, +moved to 'Starting' state, and sent to the worker. + +4. The worker creates a (paused) 'SubPubSource' for the job and notifies the +server with the 'JobStartedNotify'. The worker then waits on the result of the +'SubPubSource'. + +5. The server notifies the job originator of the 'SubPubSource' with the +'JobStartingNotify' included in the 'enqueueJob' request. + +6. The originator subscribes to the 'SubPubSource' and returns a confirmation to +the server. + +7. The server starts the 'SubPubSource', moves the job to the 'Running' state, +and waits for notification on the 'JobFinishedChan'. + +8. Optional: The server may kill the job due to a 'killJobs' request. In this +case the worker process running the job is asked to exit and the job is moved to +the 'Killed' state. + +9. The worker eventually finishes the job and notifies the server on the +'JobFinishedChan'. + +10. The server moves the job to either 'Failed' or 'Finished' state. + +-} diff --git a/TPar/SubPubStream.hs b/TPar/SubPubStream.hs new file mode 100644 --- /dev/null +++ b/TPar/SubPubStream.hs @@ -0,0 +1,178 @@ +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE BangPatterns #-} + +module TPar.SubPubStream + ( SubPubSource + , fromProducer + , fromProducer' + , subscribe + ) where + +import Control.Monad.Catch +import Control.Monad (void) +import qualified Data.Map.Strict as M +import GHC.Generics (Generic) +import Data.Binary + +import Control.Distributed.Process hiding (catch) +import Control.Distributed.Process.Serializable +import Control.Concurrent.STM +import Pipes +import TPar.Utils + +newtype SubPubSource a r = SubPubSource (SendPort (SendPort (DataMsg a r), SendPort ())) + deriving (Show, Eq, Ord, Binary) + +data DataMsg a r = More !a + | Done !r + | Failed !SubPubProducerFailed + deriving (Show, Generic) +instance (Binary a, Binary r) => Binary (DataMsg a r) + +-- | Create a new 'SubPubSource' being asynchronously fed by the given +-- 'Producer'. Exceptions thrown by the 'Producer' will be thrown to +-- subscribers. +fromProducer :: forall a r. (Serializable a, Serializable r) + => Producer a Process r + -> Process (SubPubSource a r, STM (Either SubPubProducerFailed r)) +fromProducer prod0 = do + (start, subPub, getResult) <- fromProducer' prod0 + start + return (subPub, getResult) + +-- | Create a new 'SubPubSource' being asynchronously fed by the given +-- 'Producer'. The 'Producer' will not be started until the returned 'Process' +-- is executed. Exceptions thrown by the 'Producer' will be thrown to +-- subscribers. +fromProducer' :: forall a r. (Serializable a, Serializable r) + => Producer a Process r + -> Process (Process (), SubPubSource a r, STM (Either SubPubProducerFailed r)) + -- ^ returns a 'Process' action to start source, the 'SubPubSource' + -- itself, and an 'STM' action to determine request the final + -- return value. +fromProducer' prod0 = do + dataQueue <- liftIO $ atomically $ newTBQueue 10 + (subReqSP, subReqRP) <- newChan + resultVar <- liftIO $ atomically newEmptyTMVar + feeder <- spawnLocal $ do + () <- expect + r <- feedChan dataQueue prod0 + liftIO $ atomically $ putTMVar resultVar r + + void $ spawnLocal $ do + feederRef <- monitor feeder + loop feederRef subReqRP dataQueue M.empty + + return (send feeder (), SubPubSource subReqSP, readTMVar resultVar) + where + -- Feed data from Producer into TChan + feedChan :: TBQueue (DataMsg a r) + -> Producer a Process r + -> Process (Either SubPubProducerFailed r) + feedChan queue = handleAll uhOh . go + where + uhOh exc = do + pid <- getSelfPid + let x = SubPubProducerFailed pid $ DiedException $ show exc + liftIO $ atomically $ writeTBQueue queue (Failed x) + return $ Left x + + go prod = do + mx <- next prod + case mx of + Left r -> do + tparDebug "feedChan:finishing" + liftIO $ atomically $ writeTBQueue queue (Done r) + tparDebug "feedChan:finished" + return $ Right r + + Right (x, prod') -> do + tparDebug "feedChan:fed" + x `seq` tparDebug "feedChan:forced" + liftIO $ atomically $ writeTBQueue queue (More x) + go prod' + + -- Accept requests for subscriptions and sends data downstream + loop :: MonitorRef -- ^ on the feeder + -> ReceivePort (SendPort (DataMsg a r), SendPort ()) + -- ^ where we take subscription requests + -> TBQueue (DataMsg a r) + -- ^ data from feeder + -> M.Map MonitorRef (SendPort (DataMsg a r)) + -- ^ active subscribers + -> Process () + loop feederRef subReqRP dataQueue subscribers = do + tparDebug "loop:preMatch" + receiveWait + [ -- handle death of a subscriber + matchIf (\(ProcessMonitorNotification mref _ _) -> mref `M.member` subscribers) + $ \(ProcessMonitorNotification mref _pid _reason) -> do + tparDebug "loop:subDied" + loop feederRef subReqRP dataQueue (M.delete mref subscribers) + + -- subscription request + , matchChan subReqRP $ \(sink, confirm) -> do + tparDebug "loop:subReq" + sinkRef <- monitorPort sink + sendChan confirm () + loop feederRef subReqRP dataQueue (M.insert sinkRef sink subscribers) + + -- data for subscribers + , matchSTM (readTBQueue dataQueue) $ \msg -> do + tparDebug "loop:data" + sendToSubscribers msg + case msg of + More _ -> + loop feederRef subReqRP dataQueue subscribers + _ -> return () + + -- handle death of the feeder + , matchIf (\(ProcessMonitorNotification mref _ _) -> mref == feederRef) + $ \(ProcessMonitorNotification _ pid reason) -> do + tparDebug "loop:feederDied" + sendToSubscribers $ Failed $ SubPubProducerFailed pid reason + ] + where + sendToSubscribers msg = + mapM_ (`sendChan` msg) subscribers + `catchAll` + \exc -> say $ show exc + +-- | An exception indicating that the upstream 'Producer' feeding a +-- 'SubPubSource' failed. +data SubPubProducerFailed = SubPubProducerFailed ProcessId DiedReason + deriving (Show, Generic) +instance Exception SubPubProducerFailed +instance Binary SubPubProducerFailed + +-- | Subscribe to a 'SubPubSource'. Exceptions thrown by the 'Producer' feeding +-- the 'SubPubSource' will be thrown by the returned 'Producer'. Will return +-- 'Nothing' if the 'SubPubSource' terminated before we were able to subscribe. +subscribe :: forall a r. (Serializable a, Serializable r) + => SubPubSource a r -> Process (Maybe (Producer a Process r)) +subscribe (SubPubSource reqSP) = do + -- We provide a channel to confirm that we have actually been subscribed + -- so that we can safely link during negotiation. + tparDebug "subscribing" + mref <- monitorPort reqSP + (confirmSp, confirmRp) <- newChan + (dataSp, dataRp) <- newChan + sendChan reqSP (dataSp, confirmSp) + tparDebug "subscribe: waiting for confirmation" + let go = do + msg <- lift $ receiveWait + [ matchChan dataRp return + , matchIf (\(PortMonitorNotification mref' _ _) -> mref == mref') + $ const $ fail "subscribe: it died" + ] + case msg of + Failed err -> lift $ throwM err + Done r -> return r + More x -> yield x >> go + + receiveWait + [ matchChan confirmRp $ \() -> return $ Just go + , matchIf (\(PortMonitorNotification mref' _ _) -> mref == mref') (pure $ pure Nothing) + ] diff --git a/TPar/SubPubStream/Test.hs b/TPar/SubPubStream/Test.hs new file mode 100644 --- /dev/null +++ b/TPar/SubPubStream/Test.hs @@ -0,0 +1,128 @@ +{-# LANGUAGE ScopedTypeVariables #-} + +module TPar.SubPubStream.Test where + +import Control.Exception +import Control.Applicative +import Test.QuickCheck + +import Control.Concurrent.STM +import Control.Monad.Trans.State +import Control.Distributed.Process +import Control.Distributed.Process.Node +import Control.Distributed.Process.Serializable +import Network.Transport.InMemory + +import Pipes +import qualified Pipes.Prelude as P.P +import TPar.SubPubStream + +data TestEvent a r + = NewSubscriber (Maybe Int) (TestEvent a r) -- how long before terminating + | Produce a (TestEvent a r) + | Finish r + | Throw + deriving (Show) + +instance (Arbitrary a, Arbitrary r) => Arbitrary (TestEvent a r) where + arbitrary = oneof + [ Produce <$> arbitrary <*> arbitrary + , newSub + , Finish <$> arbitrary + , pure Throw + ] + where + newSub = (\n -> NewSubscriber $ fmap getPositive n) <$> arbitrary <*> arbitrary + +atomically' :: MonadIO m => STM a -> m a +atomically' = liftIO . atomically + +test :: TestEvent Int () -> Property +test = ioProperty . runLocalProcess . test' + +--test :: Property +--test = ioProperty $ runLocalProcess $ test' $ events' + +runLocalProcess :: Process a -> IO a +runLocalProcess process = do + tport <- createTransport + node <- newLocalNode tport initRemoteTable + resultVar <- atomically newEmptyTMVar + runProcess node $ do + r <- process + say "runLocalProcess:done" + atomically' $ putTMVar resultVar r + + atomically $ takeTMVar resultVar + +produceTMVar :: TMVar (Either r a) -> Producer a Process r +produceTMVar chan = go + where + go = do + lift $ say "produceTChan:waiting" + mx <- liftIO $ atomically $ takeTMVar chan + case mx of + Right x -> lift (say "produceTChan:fed") >> yield x >> go + Left r -> lift (say "produceTChan:done") >> return r + +test' :: forall a r. (Show a, Show r, Serializable a, Serializable r) + => TestEvent a r -> Process Bool +test' events0 = do + produceChan <- atomically' newEmptyTMVar + (pubSubSrc, readRet) <- fromProducer $ produceTMVar produceChan >-> traceP + goodVars <- execStateT (go produceChan pubSubSrc events0) [] + say $ "goodVars:"++show (length goodVars) + and <$> mapM (atomically' . readTMVar) goodVars + where + go :: TMVar (Either r a) + -> SubPubSource a r + -> TestEvent a r + -> StateT [TMVar Bool] Process () + go produceChan pubSubSrc (NewSubscriber maybeN rest) = do + goodVar <- atomically' newEmptyTMVar + Just prod <- lift $ subscribe pubSubSrc + void $ lift $ spawnLocal $ do + say "testSubscriber:starting" + result <- P.P.toListM' + $ fmap Right prod + >-> fmap Left (maybe cat P.P.take maybeN) + say $ "testSubscriber:got:"++show result + say $ "testSubscriber:expected:"++show rest + -- TODO: Actually check this result + atomically' $ putTMVar goodVar True + modify (goodVar:) + go produceChan pubSubSrc rest + + go produceChan pubSubSrc (Produce x rest) = do + lift $ say "test:produce" + atomically' $ putTMVar produceChan (Right x) + go produceChan pubSubSrc rest + + go produceChan pubSubSrc (Finish x) = do + lift $ say "test:finish" + atomically' $ putTMVar produceChan (Left x) + + go produceChan pubSubSrc Throw = do + lift $ say "test:throw" + atomically' $ putTMVar produceChan (Right $ throw TestException) + +data TestException = TestException + deriving (Show) +instance Exception TestException + +events' :: TestEvent Int () +events' = NewSubscriber Nothing $ Finish () + +events :: TestEvent Int () +events = + Produce 1 + $ NewSubscriber Nothing + $ Produce 2 + $ Produce 3 + $ Finish () + +dbg :: MonadIO m => Show a => a -> m () +dbg = liftIO . print + +traceP :: (MonadIO m, Show a) => Pipe a a m r +traceP = P.P.mapM (\x -> dbg x >> return x) diff --git a/TPar/Types.hs b/TPar/Types.hs new file mode 100644 --- /dev/null +++ b/TPar/Types.hs @@ -0,0 +1,125 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE RecordWildCards #-} + +module TPar.Types where + +import Control.Applicative +import System.Exit + +import Data.Binary +import GHC.Generics +import Control.Distributed.Process +import Data.Time.Clock +import Data.Time.Calendar + +import TPar.ProcessPipe +import TPar.SubPubStream +import TPar.Rpc + +newtype JobId = JobId Int + deriving (Eq, Ord, Show, Enum, Binary) + +data Job = Job { jobId :: !JobId + , jobRequest :: JobRequest + , jobState :: !JobState + , jobStartingNotify :: Maybe JobStartingNotify + -- ^ Used to notify the originator that a + -- 'SubPubSource' is available for the job. + } + deriving (Generic) + +instance Binary Job + +newtype JobName = JobName String + deriving (Show, Eq, Ord, Binary) + +-- | A RPC call provided with an enqueue request which is called before the job +-- is started by the server. This is used to setup atomic watches. In the case +-- that the task is re-enqueued, this procedure may be called more than once. If +-- the call fails the job will be started regardless. +type JobStartingNotify = RpcSendPort (SubPubSource ProcessOutput ExitCode) () + +data JobRequest = JobRequest { jobName :: !JobName + , jobPriority :: !Priority + , jobCommand :: FilePath + , jobArgs :: [String] + , jobCwd :: FilePath + , jobSinks :: OutputStreams (Maybe FilePath) + , jobEnv :: Maybe [(String, String)] + } + deriving (Show, Generic) + +instance Binary JobRequest + +newtype Priority = Priority Int + deriving (Eq, Ord, Show, Enum, Binary) + +data JobState + = -- | the job is waiting to be run + Queued { jobQueueTime :: !UTCTime + -- ^ when was the job enqueued? + } + + -- | the job is currently starting on the given worker. + | Starting { jobProcessId :: !ProcessId + -- ^ where is the job running? + , jobQueueTime :: !UTCTime + , jobStartingTime :: !UTCTime + -- ^ when did the job begin starting? + } + -- | the job currently running on the worker with the given 'ProcessId' + | Running { jobProcessId :: !ProcessId + , jobMonitor :: !(SubPubSource ProcessOutput ExitCode) + -- ^ where can we monitor the job output? + , jobQueueTime :: !UTCTime + , jobStartTime :: !UTCTime + } + -- | the job has finished with the given 'ExitCode' + | Finished { jobExitCode :: !ExitCode + -- ^ what code did the job exit with? + , jobWorkerNode :: !NodeId + -- ^ which node was the job running on? + , jobQueueTime :: !UTCTime + , jobStartTime :: !UTCTime + , jobFinishTime :: !UTCTime + } + -- | something happened to the worker which was running the job + | Failed { jobErrorMsg :: !String + , jobWorkerNode :: !NodeId + , jobQueueTime :: !UTCTime + , jobStartTime :: !UTCTime + , jobFailedTime :: !UTCTime + } + -- | the job was manually killed (perhaps before it was even started) + | Killed { jobQueueTime :: !UTCTime + , jobKilledWorkerNode :: !(Maybe NodeId) + -- ^ which node was the job running on + , jobKilledStartTime :: !(Maybe UTCTime) + -- ^ when was the job started (if at all) + , jobKilledTime :: !UTCTime + -- ^ when was the job killed + } + deriving (Show, Generic) + +jobMaybeStartTime :: JobState -> Maybe UTCTime +jobMaybeStartTime (Queued{}) = Nothing +jobMaybeStartTime (Starting{}) = Nothing +jobMaybeStartTime (Running{..}) = Just jobStartTime +jobMaybeStartTime (Finished{..}) = Just jobStartTime +jobMaybeStartTime (Failed{..}) = Just jobStartTime +jobMaybeStartTime (Killed{..}) = jobKilledStartTime + +instance Binary JobState + +instance Binary DiffTime where + get = picosecondsToDiffTime <$> get + put = put . diffTimeToPicoseconds + +instance Binary Day where + get = ModifiedJulianDay <$> get + put = put . toModifiedJulianDay + +instance Binary UTCTime where + get = UTCTime <$> get <*> get + put (UTCTime a b) = put a >> put b diff --git a/TPar/Utils.hs b/TPar/Utils.hs new file mode 100644 --- /dev/null +++ b/TPar/Utils.hs @@ -0,0 +1,12 @@ +module TPar.Utils where + +import Control.Distributed.Process + +debugEnabled :: Bool +debugEnabled = False + +tparDebug :: String -> Process () +tparDebug _ | not debugEnabled = return () +tparDebug s = say s' + where + s' = "tpar: "++s diff --git a/Test.hs b/Test.hs new file mode 100644 --- /dev/null +++ b/Test.hs @@ -0,0 +1,5 @@ +import Test.QuickCheck +import qualified TPar.SubPubStream.Test + +main :: IO () +main = quickCheck TPar.SubPubStream.Test.test diff --git a/tpar.cabal b/tpar.cabal new file mode 100644 --- /dev/null +++ b/tpar.cabal @@ -0,0 +1,157 @@ +name: tpar +version: 0.1.0.0 +synopsis: simple, parallel job scheduling +description: + @tpar@ is a simple job scheduling and dispatch service for distributing and + monitoring tasks across machines. It was written to serve as a simple and + easy-to-administer substitute for systems like Grid Engine. + . + Configuring a @tpar@ instance is simply a matter of running @tpar server@ on + a designated server machine, + . + @ + $ # We'll need to know the hostname of the server + $ hostname + my-server + $ # Start a server also running 8 local workers + $ tpar server -Hmy-server -N8 + @ + . + Submitting a job is then similarly easy, + . + @ + $ tpar enqueue -H`hostname` -- long-process arg1 arg2 + @ + . + One can then enqueue jobs easily + . + @ + $ tpar enqueue -Hmy-server -- bash -c "primes 10000000000 | head -n1" + $ tpar status -v + 0 unnamed-job finished + priority: 0 + queued: 1 seconds ago + command: bash + arguments: -c primes 10000000000 | head -n1 + logging: stdout: none + stderr: none + status: finished with 0 (at 1 seconds ago) + started at 1 seconds ago + ran on nid://localhost:5757:0 + @ + . + One can add more workers to help churn through the work queue using the + @tpar worker@ command, + . + @ + $ # Add 16 more workers running on another machine + $ ssh my-workers -- tpar worker -Hmy-server -N16 + @ + . + Finally, the output of running jobs can be monitored using the @tpar watch@ + command, + . + @ + $ tpar watch id=0 + @ + +homepage: http://github.com/bgamari/tpar +license: BSD3 +license-file: LICENSE +author: Ben Gamari +maintainer: bgamari@gmail.com +copyright: (c) 2014-2016 Ben Gamari +category: System +build-type: Simple +homepage: http://github.com/bgamari/tpar/ +bug-reports: http://github.com/bgamari/tpar/issues +cabal-version: >=1.10 +extra-source-files: + README.mkd + +source-repository head + location: git://github.com/bgamari/tpar + type: git + +executable tpar + main-is: Main.hs + other-modules: TPar.Server + TPar.Server.Types + TPar.JobMatch + TPar.ProcessPipe + TPar.SubPubStream + TPar.Rpc + TPar.Types + TPar.Utils + other-extensions: DeriveGeneric + build-depends: base >=4.5 && <4.10, + ghc-prim, + binary >=0.7 && <0.9, + bytestring >=0.10 && <0.11, + containers >= 0.5 && < 0.6, + errors >=2.0 && <2.2, + transformers >=0.3 && <0.6, + exceptions >= 0.8 && < 0.9, + async >=2.0 && <2.2, + stm >=2.4 && <2.5, + time >=1.6 && <1.7, + friendly-time >=0.4 && <0.5, + process >=1.1 && <1.5, + pipes >=4.0 && <4.3, + pipes-bytestring >=2.0 && <2.2, + pipes-safe >= 2.2 && < 2.3, + pipes-concurrency, + optparse-applicative >= 0.10 && <0.13, + distributed-process >= 0.6 && <0.7, + network-transport-tcp >= 0.4, + network >= 2.4 && < 2.7, + heaps >= 0.3 && < 0.4, + ansi-wl-pprint >= 0.6 && < 0.7, + trifecta >= 1.5 && < 1.7, + parsers >= 0.12 && < 0.13, + aeson >= 0.11 && < 0.12 + default-language: Haskell2010 + ghc-options: -threaded -rtsopts -Wall + +test-suite test + type: exitcode-stdio-1.0 + main-is: Test.hs + other-modules: TPar.SubPubStream.Test + default-language: Haskell2010 + build-depends: base >=4.5 && <4.10, + binary, + stm, + containers, + transformers, + distributed-process, + network-transport-inmemory, + exceptions, + pipes, + QuickCheck + +benchmark bench + type: exitcode-stdio-1.0 + main-is: Benchmark.hs + build-depends: base >=4.5 && <4.10, + ghc-prim, + binary, + bytestring, + containers, + errors, + transformers, + exceptions, + async, + stm, + process, + pipes, + pipes-bytestring, + pipes-safe, + pipes-concurrency, + optparse-applicative, + distributed-process, + network-transport-tcp, + network, + heaps, + trifecta + default-language: Haskell2010 + ghc-options: -threaded -eventlog -rtsopts