Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Transfer api-key in file not env for condor jobs #227

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 46 additions & 4 deletions batchspawner/batchspawner.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,9 @@
import pwd
import os
import re
import sys
import grp

import tempfile

import xml.etree.ElementTree as ET

Expand Down Expand Up @@ -877,20 +879,26 @@ def get_env(self):
class CondorSpawner(UserEnvMixin, BatchSpawnerRegexStates):
batch_script = Unicode(
"""
Executable = /bin/sh
Executable = /bin/bash
RequestMemory = {memory}
RequestCpus = {nprocs}
Arguments = \"-c 'exec {cmd}'\"
Arguments = \"-c 'source $_CONDOR_SCRATCH_DIR/{apikey_file}; exec {cmd}'\"
Remote_Initialdir = {homedir}
Output = {homedir}/.jupyterhub.condor.out
Error = {homedir}/.jupyterhub.condor.err
ShouldTransferFiles = False
GetEnv = True
transfer_executable = False
transfer_input_files = {apikeyfile_dir}/{apikey_file}
should_transfer_files = YES
{options}
Queue
"""
).tag(config=True)

req_apikey_file = Unicode("")

req_apikeyfile_dir = Unicode("/tmp")

# outputs job id string
batch_submit_cmd = Unicode("condor_submit").tag(config=True)
# outputs job data XML string
Expand All @@ -903,6 +911,40 @@ class CondorSpawner(UserEnvMixin, BatchSpawnerRegexStates):
state_running_re = Unicode(r"^2,").tag(config=True)
state_exechost_re = Unicode(r"^\w*, .*@([^ ]*)").tag(config=True)

def write_apikey_file(self):
p = "bsapikey-{0}".format(self.user.name)

with tempfile.NamedTemporaryFile(
delete=False, dir=self.req_apikeyfile_dir, prefix=p
) as fp:
self.log.info("Writing apikey to file: %s", fp.name)
fp.write("export JUPYTERHUB_API_TOKEN={}\n".format(self.api_token).encode())
self.req_apikey_file = os.path.basename(fp.name)

# Set file owned by user for batch-submission
user = pwd.getpwnam(self.user.name)
os.chown(fp.name, user.pw_uid, user.pw_gid)

def clean_apikey_file(self):
try:
os.unlink(os.path.join(self.req_apikeyfile_dir, self.req_apikey_file))
except OSError:
pass

def get_env(self):
env = super().get_env()
env.pop("JUPYTERHUB_API_TOKEN", None)
env.pop("JPY_API_TOKEN", None)
return env

async def submit_batch_script(self):
self.write_apikey_file()
return await super().submit_batch_script()

async def cancel_batch_job(self):
self.clean_apikey_file()
await super().cancel_batch_job()

def parse_job_id(self, output):
match = re.search(r".*submitted to cluster ([0-9]+)", output)
if match:
Expand Down