Argonix.hs 6.88 KB
Newer Older
Valentin Reis's avatar
Valentin Reis committed
1 2
{-# language RecordWildCards #-}
{-# language ApplicativeDo #-}
Valentin Reis's avatar
Valentin Reis committed
3
{-# language LambdaCase #-}
Valentin Reis's avatar
Valentin Reis committed
4 5
{-# language OverloadedStrings #-}

Valentin Reis's avatar
Valentin Reis committed
6 7 8 9 10 11 12
{-|
Module      : Main
Description : argonix
Copyright   : (c) Valentin Reis, 2018
License     : MIT
Maintainer  : fre@freux.fr -}

Valentin Reis's avatar
Valentin Reis committed
13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30
module Main where

import           Data.Default
import           Turtle                  hiding ( header
                                                , text
                                                , arguments
                                                )
import           Options.Applicative
import           Data.Maybe
import           Prelude                 hiding ( FilePath )

import           Data.Text                      ( pack
                                                , unpack
                                                )
import           System.Posix.Process


data StackArgs = StackArgs
Valentin Reis's avatar
Valentin Reis committed
31 32 33
  {
    argopkgs      :: Text
   , verbosity     :: Verbosity
Valentin Reis's avatar
Valentin Reis committed
34 35 36 37 38 39 40 41 42 43 44 45
   , enableApps    :: Bool
   , remoteBuild   :: Bool
   , targetMachine :: Maybe Text
   , retreive      :: Maybe Text
   , retreiveAs    :: Maybe Text
   , run           :: Maybe Text
   , overrides     :: [(String, FilePath)]
  } deriving (Show)
data Verbosity = Verbose | Normal deriving (Show)

instance Default StackArgs where
  def = StackArgs
Valentin Reis's avatar
Valentin Reis committed
46 47
    { verbosity = Normal,
      argopkgs = "<argopkgs>",
Valentin Reis's avatar
Valentin Reis committed
48 49 50 51 52 53 54 55 56 57 58
      enableApps = False,
      remoteBuild = False,
      run = Nothing,
      targetMachine = Nothing,
      retreive = Nothing,
      retreiveAs = Nothing,
      overrides = []
    }

executorParser :: Parser StackArgs
executorParser = do
Valentin Reis's avatar
Valentin Reis committed
59 60 61 62 63 64 65
  argopkgs <- strOption
    (  long "argopkgs"
    <> metavar "ARGOPKGS"
    <> showDefault
    <> value (argopkgs def)
    <> help "Nix expression that produces the argopkgs source path."
    )
Valentin Reis's avatar
Valentin Reis committed
66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103
  verbosity <- flag
    Normal
    Verbose
    (long "verbose" <> short 'v' <> help "Enable verbose mode")
  enableApps <- flag
    False
    True
    (long "provision_apps" <> short 'a' <> help
      "Build and provision applications."
    )
  remoteBuild <- flag
    False
    True
    (long "remote_build" <> short 'r' <> help
      "Build on the remote rather than locally."
    )
  retreive <- optional $ strOption
    (long "retreive" <> metavar "RELATIVE_PATH" <> help
      "File/Folder to retreive from the remote machine."
    )
  retreiveAs <- optional $ strOption
    (long "retreiveAs" <> metavar "LOCAL_PATH" <> help
      "File/Folder to save retreived data as (locally)."
    )
  targetMachine <- optional $ strOption
    (long "target_machine" <> metavar "USER@HOST" <> help
      "Target machine. defaults to localhost via forking."
    )
  overrides <- catMaybes <$> truc
  run       <- optional $ strOption
    (long "run" <> metavar "COMMAND" <> help
      "Command to run the environment instead of an interactive shell"
    )
  pure StackArgs {..}
 where
  truc :: Parser [Maybe (String, FilePath)]
  truc = traverse
    optSrc
Valentin Reis's avatar
Valentin Reis committed
104 105
    [ "aml"
    , "libnrm"
Valentin Reis's avatar
Valentin Reis committed
106 107
    , "nrm"
    , "containers"
Valentin Reis's avatar
Valentin Reis committed
108
    , "argotk"
Valentin Reis's avatar
Valentin Reis committed
109 110 111 112
    , "amg"
    , "lammps"
    , "qmcpack"
    , "stream"
Valentin Reis's avatar
Valentin Reis committed
113
    , "openmc"
Valentin Reis's avatar
Valentin Reis committed
114 115 116 117 118 119 120 121 122 123 124 125 126 127 128
    ]
  optSrc :: String -> Parser (Maybe (String, FilePath))
  optSrc longform = do
    parsed <- optional $ strOption
      (long longform <> metavar "PATH" <> help
        (longform <> " source folder override.")
      )
    pure $ mapT longform parsed
  mapT :: String -> Maybe FilePath -> Maybe (String, FilePath)
  mapT longform thePath = case thePath of
    Nothing -> Nothing
    Just p  -> Just (longform, p)

main :: IO ()
main = sh $ do
Valentin Reis's avatar
Valentin Reis committed
129 130 131 132 133 134 135
  StackArgs {..} <- liftIO $ execParser opts

  -- building nixArguments (pure stuff) and shellArguments (impure stuff)
  let nixArguments = ["-A", if enableApps then "expe" else "test"] ++ concat
        [ ["--arg", longform <> "-src", encodeString p]
        | (longform, p) <- overrides
        ]
Valentin Reis's avatar
Valentin Reis committed
136 137 138 139
  sudo <- which "sudo" >>= \case
    (Just sudo) -> printf ("Found sudo at " % fp % "\n") sudo >> return sudo
    Nothing     -> die "sudo not in $PATH."
  export "SUDO" $ pack $ encodeString sudo
Valentin Reis's avatar
Valentin Reis committed
140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159
  let shellArguments =
        [unpack argopkgs]
          ++ [ "--keep"
             , "SUDO"
             , "--pure"
             , "--allow-new-privileges"
             , "--option"
             , "build-extra-sandbox-paths"
             , encodeString (directory sudo)
             ]
          ++ nixArguments
          ++ (case run of
               Just cmd -> ["--run", unpack ("\"exec " <> cmd <> "\"")]
               Nothing  -> []
             )

  cachedir <- single $ inproc
    "mktemp"
    ["-d", "--suffix=nixcache", "/tmp/deletable-nix-cache-XXXX"]
    empty
Valentin Reis's avatar
Valentin Reis committed
160
  export "XDG_CACHE_HOME" $ lineToText cachedir
Valentin Reis's avatar
Valentin Reis committed
161
  case targetMachine of
Valentin Reis's avatar
Valentin Reis committed
162 163 164 165
    Nothing -> do
      printf s "Running nix-shell with the following arguments: \n"
      liftIO $ print shellArguments
      liftIO $ executeFile "nix-shell" True shellArguments Nothing
Valentin Reis's avatar
Valentin Reis committed
166 167 168 169 170 171
    Just host -> do
      export "NIX_SSHOPTS" "source .profile; source .bash_profile;"
      printf s "argonix: creating a derivation in the local store:\n"

      drv <- single $ inproc
        "nix-instantiate"
Valentin Reis's avatar
Valentin Reis committed
172
        (["--quiet", argopkgs, "-A", "test"] ++ map pack nixArguments)
Valentin Reis's avatar
Valentin Reis committed
173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199
        empty

      printf (s % "\n") $ lineToText drv

      if remoteBuild
        then do
          printf
            s
            "argonix: copying the derivation's closure (without build output).\n"
          proc "nix-copy-closure" ["--to", host, lineToText drv] empty
        else do
          printf s "argonix: building the derivation's output:\n"
          proc "nix-store" ["--realize", lineToText drv, "--quiet"] empty

          printf
            s
            "argonix: copying the derivation's closure (with build output).\n"
          proc "nix-copy-closure"
               ["--include-outputs", "--to", host, lineToText drv]
               empty

      case run of
        Just cmd -> do
          _ <- proc
            "ssh"
            [ "-t"
            , host
Valentin Reis's avatar
Valentin Reis committed
200
            , "source .profile; source .bash_profile; SUDO=$(which sudo) exec nix-shell --keep SUDO --pure --allow-new-privileges --option build-extra-sandbox-paths $(which sudo | xargs dirname) "
Valentin Reis's avatar
Valentin Reis committed
201 202
            <> lineToText drv
            <> " --run "
Valentin Reis's avatar
Valentin Reis committed
203 204 205
            <> "\""
            <> cmd
            <> "\""
Valentin Reis's avatar
Valentin Reis committed
206 207 208 209 210 211 212 213 214 215 216 217 218 219
            ]
            empty
          mapM_
            (\r -> proc "scp"
                        ["-r", host <> ":" <> r, fromMaybe "." retreiveAs]
                        empty
            )
            retreive
        Nothing -> liftIO $ executeFile
          "ssh"
          True
          [ "-t"
          , unpack host
          , unpack
Valentin Reis's avatar
Valentin Reis committed
220
            ("source .profile; source .bash_profile; SUDO=$(which sudo) exec nix-shell --keep SUDO --pure --allow-new-privileges --option build-extra-sandbox-paths $(which sudo | xargs dirname) "
Valentin Reis's avatar
Valentin Reis committed
221 222 223 224 225 226 227 228 229 230 231
            <> lineToText drv
            )
          ]
          Nothing
 where
  opts = info
    (executorParser <**> helper)
    (  fullDesc
    <> progDesc "Argo environment provisioning/deployment/execution."
    <> header "argonix"
    )