Argonix.hs 6.75 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 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23
{-# language OverloadedStrings #-}

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
24 25 26
  {
    argopkgs      :: Text
   , verbosity     :: Verbosity
Valentin Reis's avatar
Valentin Reis committed
27 28 29 30 31 32 33 34 35 36 37 38
   , 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
39 40
    { verbosity = Normal,
      argopkgs = "<argopkgs>",
Valentin Reis's avatar
Valentin Reis committed
41 42 43 44 45 46 47 48 49 50 51
      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
52 53 54 55 56 57 58
  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
59 60 61 62 63 64 65 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
  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
97 98
    [ "aml"
    , "libnrm"
Valentin Reis's avatar
Valentin Reis committed
99 100
    , "nrm"
    , "containers"
Valentin Reis's avatar
Valentin Reis committed
101
    , "argotk"
Valentin Reis's avatar
Valentin Reis committed
102 103 104 105
    , "amg"
    , "lammps"
    , "qmcpack"
    , "stream"
Valentin Reis's avatar
Valentin Reis committed
106
    , "openmc"
Valentin Reis's avatar
Valentin Reis committed
107 108 109 110 111 112 113 114 115 116 117 118 119 120 121
    ]
  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
122 123 124 125 126 127 128
  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
129 130 131 132
  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
133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152
  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
153
  export "XDG_CACHE_HOME" $ lineToText cachedir
Valentin Reis's avatar
Valentin Reis committed
154
  case targetMachine of
Valentin Reis's avatar
Valentin Reis committed
155 156 157 158
    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
159 160 161 162 163 164
    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
165
        (["--quiet", argopkgs, "-A", "test"] ++ map pack nixArguments)
Valentin Reis's avatar
Valentin Reis committed
166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192
        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
193
            , "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
194 195
            <> lineToText drv
            <> " --run "
Valentin Reis's avatar
Valentin Reis committed
196 197 198
            <> "\""
            <> cmd
            <> "\""
Valentin Reis's avatar
Valentin Reis committed
199 200 201 202 203 204 205 206 207 208 209 210 211 212
            ]
            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
213
            ("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
214 215 216 217 218 219 220 221 222 223 224
            <> lineToText drv
            )
          ]
          Nothing
 where
  opts = info
    (executorParser <**> helper)
    (  fullDesc
    <> progDesc "Argo environment provisioning/deployment/execution."
    <> header "argonix"
    )