diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -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 diff --git a/examples/Main.hs b/examples/Main.hs --- a/examples/Main.hs +++ b/examples/Main.hs @@ -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 diff --git a/lib/Neptune/AbortHandler.hs b/lib/Neptune/AbortHandler.hs new file mode 100644 --- /dev/null +++ b/lib/Neptune/AbortHandler.hs @@ -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 () + diff --git a/lib/Neptune/Channel.hs b/lib/Neptune/Channel.hs --- a/lib/Neptune/Channel.hs +++ b/lib/Neptune/Channel.hs @@ -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 diff --git a/lib/Neptune/Client.hs b/lib/Neptune/Client.hs --- a/lib/Neptune/Client.hs +++ b/lib/Neptune/Client.hs @@ -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 -- ^ \\/\ -> (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 -- ^ \\/\ + -> 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 () diff --git a/lib/Neptune/OAuth.hs b/lib/Neptune/OAuth.hs --- a/lib/Neptune/OAuth.hs +++ b/lib/Neptune/OAuth.hs @@ -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 diff --git a/lib/Neptune/Session.hs b/lib/Neptune/Session.hs --- a/lib/Neptune/Session.hs +++ b/lib/Neptune/Session.hs @@ -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 diff --git a/neptune-backend.cabal b/neptune-backend.cabal --- a/neptune-backend.cabal +++ b/neptune-backend.cabal @@ -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