neptune-backend 0.1.2 → 0.3.0
raw patch · 8 files changed
Files
- README.md +0/−1
- examples/Main.hs +1/−1
- lib/Neptune/AbortHandler.hs +78/−0
- lib/Neptune/Channel.hs +18/−5
- lib/Neptune/Client.hs +66/−21
- lib/Neptune/OAuth.hs +1/−1
- lib/Neptune/Session.hs +9/−0
- neptune-backend.cabal +15/−10
README.md view
@@ -13,7 +13,6 @@ ## TODOs * Support logging text, image value. * Create project-* Register callback for the abort action from web console. * Support system channels * Get git information
examples/Main.hs view
@@ -16,6 +16,6 @@ main = do withNept "jiasen/sandbox" $ \_ experiment -> do- forM_ [1..10::Int] $ \i -> do+ forM_ [1..200::Int] $ \i -> do nlog experiment "counter" (fromIntegral (i * i) :: Double) threadDelay 1000000
+ lib/Neptune/AbortHandler.hs view
@@ -0,0 +1,78 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE QuasiQuotes #-}+module Neptune.AbortHandler where++import Control.Lens (_Right, (^?!))+import Control.Retry (constantDelay, recoverAll)+import Data.Aeson (FromJSON (..), Value (..),+ eitherDecode', (.:))+import Data.Aeson.Types (prependFailure, typeMismatch)+import Data.Text.Encoding (encodeUtf8)+import qualified Network.WebSockets as WS+import RIO+import RIO.Text as T (pack, unpack)+import Text.URI (mkURI)+import Text.URI.Lens+import qualified Wuss++import Neptune.Backend.Model (ExperimentId (..))+import Neptune.OAuth (oas_access_token)+import Neptune.Session (Experiment, NeptuneSession, ct_api_url,+ exp_experiment_id, neptune_client_token,+ neptune_oauth2)++abortListener :: NeptuneSession -> Experiment -> ThreadId -> IO ()+abortListener sess exp main_thread = do+ base_url <- mkURI (sess ^. neptune_client_token . ct_api_url)++ let host = base_url ^?! uriAuthority . _Right . authHost . unRText+ path = notif_path <> exp_id <> "/operations" :: String++ oauth_ref = sess ^. neptune_oauth2++ run_listener = do+ oauth_current <- readMVar oauth_ref+ let oauth_token = oauth_current ^. oas_access_token+ Wuss.runSecureClientWith+ (T.unpack host) port path+ WS.defaultConnectionOptions+ [("Authorization", "Bearer " <> encodeUtf8 oauth_token)]+ (listener main_thread)++ -- usually it is the server that close the socket, however the program+ -- can run really a long time and a network error might happen.+ -- If exception (synchronous) happens, delay 0.5s and reconnect.+ recoverAll (constantDelay 500000) $ \_ -> run_listener++ where+ port = 443+ notif_path = "/api/notifications/v1/experiments/"+ exp_id = T.unpack $ unExperimentId $ exp ^. exp_experiment_id+++data Message = MessageAbort { _msg_experiment_id :: Text }++instance FromJSON Message where+ parseJSON (Object v) = do+ typ <- v .: "messageType"+ case typ of+ "Abort" -> do+ body <- v .: "messageBody"+ MessageAbort <$> body .: "experimentId"+ _ -> fail $ "Unsupported message type" <> (T.unpack typ)+ parseJSON invalid =+ prependFailure "parsing Message failed, "+ (typeMismatch "Object" invalid)++data AbortException = AbortException+ deriving Show++instance Exception AbortException++listener :: ThreadId -> WS.ClientApp ()+listener main_thread conn = forever $ do+ msg <- WS.receiveData conn+ case eitherDecode' msg of+ Right (MessageAbort _) -> throwTo main_thread AbortException+ Left msg -> traceM (T.pack msg) >> return ()+
lib/Neptune/Channel.hs view
@@ -11,11 +11,16 @@ import Control.Concurrent.Event as E (isSet, set) import Control.Lens+import Control.Retry (recoverAll, retryPolicyDefault) import Data.Typeable (Proxy (..), cast)-import RIO hiding (Lens', (^.))+import RIO hiding (Lens', (.~), (^.), (^..),+ (^?)) import qualified RIO.HashMap as M import qualified Neptune.Backend.API as NBAPI+import Neptune.Backend.Core (configLogContext,+ configLogExecWithContext)+import Neptune.Backend.Logging (_log, levelError) import Neptune.Backend.MimeTypes import Neptune.Backend.Model hiding (Experiment) import Neptune.Backend.ModelLens@@ -36,7 +41,11 @@ if stop_flag then readTChanFull _exp_outbound_q else readTChanAtMost 100 _exp_outbound_q- send dat+ -- Retry sending the whole batch upto 6 times.+ -- Discard and log if fails continuously+ handle onErr $+ recoverAll retryPolicyDefault $ \_ ->+ send dat if not stop_flag then go else E.set _exp_transmitter_flag@@ -62,12 +71,13 @@ case chn of DataChannelAny chn -> do let (errs, grouped) = gatherDataPoints (proxy chn) dat- -- TODO log properly- mapM_ print errs+ mapM_ (logE _neptune_config) errs return $ DataChannelWithData (chn, grouped) sendChannel session _exp_experiment_id chn_with_dat + onErr exc = logE _neptune_config $ "Failed to send data points.\n" <> tshow (exc :: SomeException)+ proxy :: f a -> Proxy a proxy _ = Proxy @@ -119,7 +129,7 @@ xs = err ^. batchChannelValueErrorDTOXL ecode = err ^. batchChannelValueErrorDTOErrorL . errorCodeL emsg = err ^. batchChannelValueErrorDTOErrorL . errorMessageL- print (chn, xs, ecode, emsg)+ logE _neptune_config $ tshow (chn, xs, ecode, emsg) where toChannelsValues :: DataChannelWithData -> InputChannelValues@@ -155,3 +165,6 @@ castData (DataPointAny d) = case cast d of Nothing -> Left $ _dpt_name d <> " is not compatible for the channel" Just o -> Right o++logE config msg = configLogExecWithContext config (configLogContext config) $+ _log "Neptune.Channel" levelError msg
lib/Neptune/Client.hs view
@@ -10,19 +10,27 @@ import Control.Concurrent (forkIO, killThread) import Control.Concurrent.Event as E (new, set, waitTimeout)-import Control.Lens ((<&>), (^.))+import Control.Exception (AsyncException (UserInterrupt),+ asyncExceptionFromException, try)+import Control.Lens (bimapping, each, filtered, (<&>),+ (^.), (^..)) import qualified Data.Text.Lazy as TL import qualified Data.Text.Lazy.Encoding as TL+import Data.Text.Lens (packed) import Data.Time.Clock (getCurrentTime) import qualified Data.UUID as UUID (toText) import Data.UUID.V4 as UUID (nextRandom) import qualified Network.HTTP.Client as NH import qualified Network.HTTP.Client.TLS as NH-import RIO hiding (Lens', (^.))+import RIO hiding (Lens', try, (^.), (^..)) import qualified RIO.HashMap as M import qualified RIO.Text as T+import System.Environment (getArgs, getEnvironment) import System.Envy (decodeEnv)+import System.Posix.Signals (Handler (Catch), installHandler,+ keyboardSignal) +import Neptune.AbortHandler (AbortException (..), abortListener) import qualified Neptune.Backend.API as NBAPI import Neptune.Backend.Client import Neptune.Backend.Core@@ -65,16 +73,22 @@ params "command" -- legacy (fromMaybe "Untitled" name)- tags){ experimentCreationParamsDescription = description }+ tags){ experimentCreationParamsDescription = description+ , experimentCreationParamsAbortable = Just True }+ let exp_id = ExperimentId (exp ^. experimentIdL) chan <- newTChanIO user_channels <- newTVarIO M.empty stop_flag <- E.new transmitter_flag <- E.new- let exp = Experiment exp_id chan user_channels stop_flag transmitter_flag undefined+ let exp = Experiment exp_id chan user_channels stop_flag transmitter_flag undefined undefined transmitter_thread <- forkIO $ transmitter session exp- return exp {_exp_transmitter = transmitter_thread} + parent_thread <- myThreadId+ abort_handler <- forkIO $ abortListener session exp parent_thread++ return exp {_exp_transmitter = transmitter_thread, _exp_abort_handler = abort_handler}+ where _mkParameter (ExperimentParamS name value) = do _id <- UUID.toText <$> UUID.nextRandom@@ -95,21 +109,53 @@ dat = DataPointAny $ DataPoint name now value atomically $ writeTChan chan dat --- | Run an action within a neptune session+-- | Run an action within a neptune session and a new experiment withNept :: Text -- ^ \<namespace\>\/\<project_name\> -> (NeptuneSession -> Experiment -> IO a) -- ^ action -> IO a withNept project_qualified_name act = do+ args <- T.unwords . map T.pack <$> getArgs+ envs <- getEnvironment+ let valid_pat name = T.isPrefixOf "MXNET_" name ||+ T.isPrefixOf "NVIDIA_" name ||+ T.isPrefixOf "CUDA_" name+ envs <- pure $ envs ^.. each . bimapping packed packed . filtered (valid_pat . fst)+ withNept' project_qualified_name Nothing Nothing [] (("args", args) : envs) [] act++-- | Run an action within a neptune session and a new experiment+withNept' :: Text -- ^ \<namespace\>\/\<project_name\>+ -> Maybe Text -- ^ Optional name of the experiment (automatically assigned if Nothing)+ -> Maybe Text -- ^ Optional description of the experiment+ -> [Parameter] -- ^ experiment hyper-parameters+ -> [(Text, Text)] -- ^ experiment properties+ -> [Text] -- ^ experiment tags+ -> (NeptuneSession -> Experiment -> IO a) -- ^ action+ -> IO a+withNept' project_qualified_name name description params props tags act = do ses <- initNept project_qualified_name- exp <- createExperiment ses Nothing Nothing [] [] []+ exp <- createExperiment ses name description params props tags + -- install an signal handler for CTRL-C, ensuring that an async-+ -- exception UserInterrupt is sent to the main thread+ main_thread <- myThreadId+ let interrupted = throwTo main_thread UserInterrupt++ old_handler <- installHandler keyboardSignal (Catch interrupted) Nothing result <- try (act ses exp)+ _ <- installHandler keyboardSignal old_handler Nothing+ case result of Left (e :: SomeException) -> do- teardownNept ses exp ExperimentState'Failed (T.pack $ displayException e)+ let end_state = case fromException e of+ Just AbortException -> Nothing+ _ -> Just $+ case asyncExceptionFromException e of+ Just UserInterrupt -> (ExperimentState'Failed, "User interrupted.")+ _ -> (ExperimentState'Failed, T.pack $ displayException e)+ teardownNept ses exp end_state throwM e Right a -> do- teardownNept ses exp ExperimentState'Succeeded ""+ teardownNept ses exp (Just (ExperimentState'Succeeded, "")) return a -- | Initialize a neptune session@@ -120,7 +166,7 @@ ct@ClientToken{..} <- decodeEnv >>= either throwString return mgr <- NH.newManager NH.tlsManagerSettings- config0 <- withStderrLogging =<< newConfig+ config0 <- pure . withNoLogging =<< newConfig let api_endpoint = TL.encodeUtf8 (TL.fromStrict _ct_api_url) config = config0 { configHost = api_endpoint }@@ -154,10 +200,9 @@ -- | Teardown a neptune session teardownNept :: NeptuneSession -- ^ session -> Experiment -- ^ experiment- -> ExperimentState -- ^ completion state- -> Text -- ^ completion message+ -> Maybe (ExperimentState, Text) -- ^ completion state & message -> IO ()-teardownNept NeptuneSession{..} experiment state msg = do+teardownNept NeptuneSession{..} experiment state_msg = do E.set (experiment ^. exp_stop_flag) -- wait at most 5 seconds done <- E.waitTimeout (experiment ^. exp_transmitter_flag) 5000000@@ -166,11 +211,11 @@ killThread $ experiment ^. exp_transmitter killThread $ _neptune_oauth2_refresh - _ <- _neptune_dispatch $ NBAPI.markExperimentCompleted- (ContentType MimeJSON)- (Accept MimeNoContent)- (mkCompletedExperimentParams state msg)- (experiment ^. exp_experiment_id) :: IO NoContent-- return ()-+ case state_msg of+ Just (state, msg) ->+ void (_neptune_dispatch $ NBAPI.markExperimentCompleted+ (ContentType MimeJSON)+ (Accept MimeNoContent)+ (mkCompletedExperimentParams state msg)+ (experiment ^. exp_experiment_id) :: IO NoContent)+ Nothing -> return ()
lib/Neptune/OAuth.hs view
@@ -17,7 +17,7 @@ ReqBodyUrlEnc (..), defaultHttpConfig, jsonResponse, req, responseBody, runReq, useHttpsURI, (=:))-import RIO hiding (Lens', (^.))+import RIO hiding (Lens', (.~), (^.)) import qualified RIO.Text as T import Text.URI (mkURI) import qualified Web.JWT as JWT
lib/Neptune/Session.hs view
@@ -77,6 +77,7 @@ , _exp_stop_flag :: E.Event -- ^ Event flag to indicate the end of the session , _exp_transmitter_flag :: E.Event -- ^ Event flag to indicate completion of transmission , _exp_transmitter :: ThreadId -- ^ Background thread for transmission+ , _exp_abort_handler :: ThreadId } class (Typeable a, Show a) => NeptDataType a where@@ -106,6 +107,14 @@ data DataPointAny = forall a . NeptDataType a => DataPointAny (DataPoint a) deriving instance Show DataPointAny +makeLenses ''ClientToken+makeLensesFor [("_neptune_http_manager", "neptune_http_manager")+ ,("_neptune_client_token", "neptune_client_token")+ ,("_neptune_config", "neptune_config")+ ,("_neptune_oauth2", "neptune_oauth2")+ ,("_neptune_oauth2_refresh", "neptune_oauth2_refresh")+ ,("_neptune_project", "neptune_project")]+ ''NeptuneSession makeLenses ''Experiment makeLenses ''DataPoint
neptune-backend.cabal view
@@ -1,6 +1,6 @@ cabal-version: 2.2 name: neptune-backend-version: 0.1.2+version: 0.3.0 synopsis: Neptune Client description: . Client library for calling the Neptune API based on http-client.@@ -54,15 +54,19 @@ , transformers >=0.4.0.0 , unordered-containers , vector >=0.10.9 && <0.13- , lens- , lens-aeson- , rio- , modern-uri- , jwt- , req- , envy- , uuid- , concurrent-extra+ , lens >=4 && < 5+ , lens-aeson >= 1 && < 2+ , rio >= 0.1 && < 0.2+ , modern-uri >= 0.3 && < 0.4+ , jwt < 1+ , req >= 3 && < 4+ , envy >= 2.0 && < 2.1+ , uuid >= 1 && < 2+ , concurrent-extra >= 0.7 && < 0.8+ , websockets >= 0.12 && < 0.13+ , wuss >= 1 && < 2+ , unix >= 2.7 && < 2.8+ , retry >= 0.8 && < 0.9 other-modules: Paths_neptune_backend autogen-modules:@@ -74,6 +78,7 @@ Neptune.OAuth Neptune.Session Neptune.Utils+ Neptune.AbortHandler Neptune.Backend.API Neptune.Backend.API.ApiDefault Neptune.Backend.Client