never executed always true always false
1 {- |
2 Module : NITTA.Synthesis.MlBackend.ServerInstance
3 Description : ML backend server instance management
4 Copyright : (c) Ilya Burakov, 2023
5 License : BSD3
6 Maintainer : aleksandr.penskoi@gmail.com
7 Stability : experimental
8 -}
9 module NITTA.Synthesis.MlBackend.ServerInstance (
10 MlBackendServer (..),
11 withLazyMlBackendServer,
12 ) where
13
14 import Control.Concurrent
15 import Control.Exception
16 import Control.Monad
17 import Data.Default
18 import Data.Functor ((<&>))
19 import Data.Maybe
20 import Data.String
21 import Data.String.ToString
22 import Data.Text qualified as T
23 import GHC.IO.Handle
24 import NITTA.Synthesis.MlBackend.FixedCache
25 import NITTA.Utils.Base
26 import System.Directory
27 import System.IO
28 import System.Log.Logger
29 import System.Process
30
31 data MlBackendServer = MlBackendServer
32 { baseUrl :: Maybe T.Text
33 , handles :: Maybe (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
34 }
35
36 instance Default MlBackendServer where
37 def = MlBackendServer Nothing Nothing
38
39 -- | Reads ML backend server base URL from file (non-blocking).
40 readMlBackendBaseUrlFileIO :: IO (Maybe T.Text)
41 readMlBackendBaseUrlFileIO =
42 do
43 let baseUrlFilePath = ".ml_backend_base_url"
44
45 debugM "NITTA.Synthesis.MlBackend" "reading ML backend base URL from file"
46 withFile
47 baseUrlFilePath
48 ReadMode
49 ( \file -> do
50 baseUrl <- hGetLine file
51 return $ Just $ fromString baseUrl
52 )
53 `catch` \(_ :: IOException) -> do
54 debugM "NITTA.Synthesis.MlBackend" "failed to read ML backend base URL from file"
55 return Nothing
56
57 -- | Synchonously waits until ML backend server base URL file becomes available (with timeout)
58 waitForMlBackendBaseUrlIO retriesLeft
59 | retriesLeft <= 0 = return Nothing
60 | otherwise = do
61 let secondsPerRetry = 3 :: Int
62
63 debugM
64 "NITTA.Synthesis.MlBackend"
65 ( "waiting "
66 <> show (secondsPerRetry * retriesLeft)
67 <> " second(s) more until ML backend server base URL is available..."
68 )
69
70 threadDelay $ secondsPerRetry * 1000000
71
72 readMlBackendBaseUrlFileIO >>= \case
73 Just baseUrl -> do
74 threadDelay 5000000 -- wait 5 more seconds to ensure that server is ready to accept connections
75 return $ Just baseUrl
76 Nothing -> waitForMlBackendBaseUrlIO (retriesLeft - 1)
77
78 findPythonExecutableIO = do
79 (mapM findExecutable ["python3", "python"] <&> catMaybes) >>= \case
80 executable : _ -> return executable
81 [] -> do
82 let errorMsg = "failed to find a python executable"
83 errorM "NITTA.Synthesis.MlBackend" errorMsg
84 throw $ userError errorMsg
85
86 -- | Tries to start ML backend server and gathers all required information about it.
87 tryStartMlBackendServerIO = do
88 -- not a shell process, because killing ML backend under shell is trickier
89 -- let's find out python3 executable path
90
91 maybeHandles <- catchToMaybeIO $ do
92 executable <- findPythonExecutableIO
93 let args = ["-m", "mlbackend"]
94 infoM "NITTA.Synthesis.MlBackend" $ "Starting ML backend server, executable: " <> executable <> ", args: " <> show args
95 createProcess (proc executable args)
96
97 maybeDynamicBaseUrl <- case maybeHandles of
98 Just _ -> waitForMlBackendBaseUrlIO 10
99 Nothing -> return Nothing
100
101 when (isNothing maybeDynamicBaseUrl) $ do
102 errorM "NITTA.Synthesis.MlBackend" "failed to start ML backend server"
103
104 return MlBackendServer{baseUrl = maybeDynamicBaseUrl, handles = maybeHandles}
105
106 -- | Makes ML backend server available with lazy initialization, memoization and proper cleanup.
107 withLazyMlBackendServer action =
108 bracket
109 -- resource initialization action, produces lazy server getter with enabled memoization
110 ( do
111 let startupAction = do
112 readMlBackendBaseUrlFileIO >>= \case
113 -- if ML backend server base URL file already exists, then we assume that server is already running
114 Just existingBaseUrl -> do
115 debugM
116 "NITTA.Synthesis.MlBackend"
117 ( "ML backend server base URL was found ("
118 <> toString existingBaseUrl
119 <> "), skipping server startup"
120 )
121 return MlBackendServer{baseUrl = Just existingBaseUrl, handles = Nothing}
122 -- coulnd't find existing base URL, trying to start the server
123 Nothing -> tryStartMlBackendServerIO
124 -- cache is used to memoize server startup results and produce lazy getter
125 cache <- newCache
126 let serverGetter = fetch cache startupAction
127 return (serverGetter, cache)
128 )
129 -- resource cleanup action
130 ( \(_, Cache mServerVar) -> do
131 tryReadMVar mServerVar >>= \case
132 -- server was started
133 Just (Just MlBackendServer{handles = Just serverProcessHandles@(_, _, _, procHandle)}) -> do
134 infoM "NITTA.Synthesis.MlBackend" "Stopping automatically started ML backend server"
135 cleanupProcess serverProcessHandles
136 _ <- waitForProcess procHandle
137 return ()
138 -- server was not started
139 _ -> debugM "NITTA.Synthesis.MlBackend" "ML backend server was not started automatically, so nothing to stop"
140 )
141 -- resource usage action (not exposing Cache object to it)
142 (\(serverGetter, _) -> action serverGetter)