Skip to content

Commit

Permalink
adapting job preparation to config file usage
Browse files Browse the repository at this point in the history
  • Loading branch information
telliere committed Mar 21, 2024
1 parent a905eb0 commit a91128a
Show file tree
Hide file tree
Showing 8 changed files with 35 additions and 14 deletions.
4 changes: 2 additions & 2 deletions client/container_preparation/input_logic/run.sh
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,10 @@ PATH="$PATH:/sd-container/tools/input_logic/"
echo "[SD-Container][Input-Logic] : Getting data decryption key from vault"

# Get token via vault login. The data_login environment variable need to be exported from calling script
data_token=$(curl -s --request POST --data "$data_login" http://${vault}/v1/auth/jwt/login | jq '.auth.client_token' -r) || exit 1
data_token=$(curl -s --request POST --data "$data_login" $vault/v1/auth/jwt/login | jq '.auth.client_token' -r) || exit 1

# Use the token to access the key. The data_path environment variable needs to be exported from calling script
data_key=$(curl -s -H "X-Vault-Token: $data_token" http://${vault}/v1/kv/data/${data_path} | jq '.data.data.key' -r) || exit 1
data_key=$(curl -s -H "X-Vault-Token: $data_token" $vault/v1/kv/data/${data_path} | jq '.data.data.key' -r) || exit 1

# Write the key in an encrypted volume
echo "$data_key" > /sd-container/encrypted/decryption_key
Expand Down
3 changes: 2 additions & 1 deletion client/job_preparation/lib/sbatch_generation.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,8 @@ def boostrap_from_template(options: argparse.Namespace, template_path: str) -> s
sbatch = sbatch.replace("ACCOUNT", options.account)
sbatch = sbatch.replace("NODELIST", options.nodelist)
sbatch = sbatch.replace("WORKDIR", options.workdir)
sbatch = sbatch.replace("TRUST_DOMAIN", "lumi-sd-dev")
sbatch = sbatch.replace("TRUST_DOMAIN", options.trust_domain)
sbatch = sbatch.replace("VAULT_ADDRESS", options.vault_address)

# Dataset info
sbatch = sbatch.replace("DATA_PATH", options.data_path_at_rest)
Expand Down
12 changes: 12 additions & 0 deletions client/job_preparation/prepare_job.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,22 @@
sys.path.append(os.path.expanduser("../../../")) # For cli usage
sys.path.append(os.path.expanduser("../../")) # For inside-container usage
from utils.ssh_utils import ssh_connect, ssh_copy_file, ssh_run_command
from utils.conf.client.conf import parse_configuration
from time import sleep
from pyrage import x25519

if __name__ == "__main__":
# Parse arguments
options = check_arguments(parse_arguments())

# Parse configuration
configuration = parse_configuration(options.config)

# Parse configuration as options
options.username = configuration['supercomputer']['username']
options.trust_domain = configuration['spire-server']['trust-domain']
options.vault_address = configuration['vault']['url']

# Connect via SSH to supercomputer
ssh_client = ssh_connect(options.username)

Expand All @@ -39,6 +48,9 @@

# Copy SBATCH to supercomputer
ssh_copy_file(ssh_client, sbatch_path, f"~/")

# Copy config file to supercomputer
ssh_copy_file(ssh_client, options.config, f"~/.config/hpcs.conf")

# Create public encryption key for output data
ident = x25519.Identity.generate()
Expand Down
8 changes: 4 additions & 4 deletions client/job_preparation/utils/cli/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,11 +11,11 @@ def parse_arguments() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="CLI Optinons")

parser.add_argument(
"--username",
"-u",
required=True,
"--config",
type=str,
help="username on supercomputer",
required=True,
default="/tmp/hpcs-client.conf",
help="Configuration file (INI Format) (default: /tmp/hpcs-client.conf)",
)
parser.add_argument(
"--job-name",
Expand Down
10 changes: 5 additions & 5 deletions client/job_preparation/utils/sbatch.template
Original file line number Diff line number Diff line change
Expand Up @@ -61,8 +61,8 @@ mkdir -p ${WORKING_DIRECTORY}
echo -e "${YELLOW}[LUMI-SD]${NC}${BLUE}[Job]${NC} Running agent registration"

# Spawn spire-agent
cd ~/LUMI-secure-processing || exit 1
python3 ./utils/spawn_agent.py -cn > $WORKING_DIRECTORY/agent.log 2> $WORKING_DIRECTORY/agent.log &
cd ~/HPCS || exit 1
python3 ./utils/spawn_agent.py --config ~/.config/hpcs-client.conf -cn > $WORKING_DIRECTORY/agent.log 2> $WORKING_DIRECTORY/agent.log &
spire_agent_pid=$!

# Wait until agent runs properly
Expand All @@ -85,18 +85,18 @@ echo "Logging in to the vault ..."

# Log in to the vault using SVID, access role
echo "{\"role\": \"APPLICATION_ACCESS_ROLE\", \"jwt\" : \"$svid\"}" > /tmp/login
application_token=$(curl -s --request POST --data @/tmp/login http://${vault}/v1/auth/jwt/login | jq '.auth.client_token' -r) || cleanup $spire_agent_pid 1
application_token=$(curl -s --request POST --data @/tmp/login $vault/v1/auth/jwt/login | jq '.auth.client_token' -r) || cleanup $spire_agent_pid 1

echo "Getting container decryption key ..."

# Use provided vault token (from login) to access secrets
data_key=$(curl -s -H "X-Vault-Token: $application_token" http://${vault}/v1/kv/data/APPLICATION_SECRET_PATH | jq '.data.data.key' -r) || cleanup $spire_agent_pid 1
data_key=$(curl -s -H "X-Vault-Token: $application_token" $vault/v1/kv/data/APPLICATION_SECRET_PATH | jq '.data.data.key' -r) || cleanup $spire_agent_pid 1
echo "$data_key" > /tmp/container_key

echo "Decrypting container image ..."

# Decrypt the container image
~/LUMI-secure-processing/client/container_preparation/input_logic/age --decrypt -i /tmp/container_key -o $WORKING_DIRECTORY/app.sif APPLICATION_PATH || exit 1
~/HPCS/client/container_preparation/input_logic/age --decrypt -i /tmp/container_key -o $WORKING_DIRECTORY/app.sif APPLICATION_PATH || exit 1

echo -e "${YELLOW}[LUMI-SD]${NC}${BLUE}[Job]${NC} Creating encrypted volumes"

Expand Down
6 changes: 6 additions & 0 deletions utils/conf/client/conf.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,9 @@ def parse_configuration(path : str):
config = ConfigParser()
config.read(path)

if not 'supercomputer' in config:
raise NoSectionError("supercomputer section missing in configuration file, aborting")

if not 'spire-server' in config:
raise NoSectionError("hpcs-server section missing in configuration file, aborting")

Expand All @@ -14,6 +17,9 @@ def parse_configuration(path : str):
if not 'vault' in config:
raise NoSectionError("vault section missing in configuration file, aborting")

if not 'address' in config['supercomputer'] or not 'username' in config['supercomputer']:
raise NoOptionError("'spire-server' section is incomplete in configuration file, aborting")

if not 'address' in config['spire-server'] or not 'port' in config['spire-server'] or not 'trust-domain' in config['spire-server']:
raise NoOptionError("'spire-server' section is incomplete in configuration file, aborting")

Expand Down
3 changes: 2 additions & 1 deletion utils/ship_a_key.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,8 @@ def parse_arguments() -> argparse.ArgumentParser:
parser.add_argument(
"--config",
required=True,
help="Path to the client configuration file",
default="/tmp/hpcs-client.conf",
help="Configuration file (INI Format) (default: /tmp/hpcs-client.conf)",
)
parser.add_argument(
"--users",
Expand Down
3 changes: 2 additions & 1 deletion utils/spawn_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,8 @@ def parse_arguments():
parser.add_argument(
"--config",
required=True,
help="Path to the client configuration file",
default="/tmp/hpcs-client.conf",
help="Configuration file (INI Format) (default: /tmp/hpcs-client.conf)",
)
parser.add_argument(
"--socketpath",
Expand Down

0 comments on commit a91128a

Please sign in to comment.