Argonix.hs 5.95 KB
Newer Older
Valentin Reis's avatar
Valentin Reis committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 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 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114
{-# language RecordWildCards #-}
{-# language ApplicativeDo #-}
{-# language OverloadedStrings #-}

module Main where

import           Data.Default
import           Nixexpr
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
  {  verbosity     :: Verbosity
   , 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
    { verbosity  = Normal,
      enableApps = False,
      remoteBuild = False,
      run = Nothing,
      targetMachine = Nothing,
      retreive = Nothing,
      retreiveAs = Nothing,
      overrides = []
    }

executorParser :: Parser StackArgs
executorParser = do
  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
    [ "argotk"
    , "argopkgs"
    , "nrm"
    , "containers"
    , "libnrm"
    , "amg"
    , "lammps"
    , "qmcpack"
    , "stream"
    ]
  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
  a@StackArgs {..} <- liftIO $ execParser opts
  let arguments = argumentsBuilder a
  case targetMachine of
    Nothing -> liftIO $ executeFile
115
      "SUDO=$(which sudo) nix-shell --keep SUDO --pure --allow-new-privileges --option build-extra-sandbox-paths $(which sudo | xargs dirname)"
Valentin Reis's avatar
Valentin Reis committed
116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158
      True
      (  ["-E", unpack nixExpression]
      ++ arguments
      ++ (case run of
           Just cmd -> ["--run", unpack ("\"" <> cmd <> "\"")]
           Nothing  -> []
         )
      )
      Nothing
    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"
        (["--quiet", "-E", nixExpression] ++ map pack arguments)
        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
159
            , "source .profile; source .bash_profile; SUDO=$(which sudo) nix-shell --keep SUDO --pure --allow-new-privileges --option build-extra-sandbox-paths $(which sudo | xargs dirname) "
Valentin Reis's avatar
Valentin Reis committed
160 161
            <> lineToText drv
            <> " --run "
Valentin Reis's avatar
Valentin Reis committed
162 163 164
            <> "\""
            <> cmd
            <> "\""
Valentin Reis's avatar
Valentin Reis committed
165 166 167 168 169 170 171 172 173 174 175 176 177 178
            ]
            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
179
            ("source .profile; source .bash_profile; SUDO=$(which sudo) nix-shell --keep SUDO --pure --allow-new-privileges --option build-extra-sandbox-paths $(which sudo | xargs dirname) "
Valentin Reis's avatar
Valentin Reis committed
180 181 182 183 184 185 186 187
            <> lineToText drv
            )
          ]
          Nothing
 where
  argumentsBuilder :: StackArgs -> [String]
  argumentsBuilder StackArgs {..} =
    ["-A", if enableApps then "expe" else "test"] ++ concat
Valentin Reis's avatar
Valentin Reis committed
188 189 190 191
      [ [ "--arg"
        , longform <> "-src"
        , unpack $ filterSource (pack $ encodeString p)
        ]
Valentin Reis's avatar
Valentin Reis committed
192 193 194 195 196 197 198 199 200
      | (longform, p) <- overrides
      ]

  opts = info
    (executorParser <**> helper)
    (  fullDesc
    <> progDesc "Argo environment provisioning/deployment/execution."
    <> header "argonix"
    )