Merge 016b6f5 into b21d474

getzlab · May 1, 2020 · 184217b · 184217b
2 parents b21d474 + 016b6f5
commit 184217b
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 9 deletions.
diff --git a/canine/localization/base.py b/canine/localization/base.py
@@ -517,27 +517,33 @@ def job_setup_teardown(self, jobId: str, patterns: typing.Dict[str, str]) -> typ
 		# - localization tasks are run when localization.sh is _run_
         job_vars = []
         exports = []
+        docker_args = ['-v $CANINE_ROOT:$CANINE_ROOT']
         localization_tasks = [
             'if [[ -d $CANINE_JOB_INPUTS ]]; then cd $CANINE_JOB_INPUTS; fi'
         ]
         compute_env = self.environment('remote')
+        stream_dir_ready = False
         for key, val in self.inputs[jobId].items():
             if val.type == 'stream':
                 job_vars.append(shlex.quote(key))
-                dest = self.reserve_path('jobs', jobId, 'inputs', os.path.basename(os.path.abspath(val.path)))
+                if not stream_dir_ready:
+                    exports.append('export CANINE_STREAM_DIR=$(mktemp -d /tmp/canine_streams.$SLURM_ARRAY_JOB_ID.$SLURM_ARRAY_TASK_ID.XXXX)')
+                    docker_args.append('-v $CANINE_STREAM_DIR:$CANINE_STREAM_DIR')
+                    stream_dir_ready = True
+                dest = os.path.join('$CANINE_STREAM_DIR', os.path.basename(os.path.abspath(val.path)))
                 localization_tasks += [
                     'gsutil ls {} > /dev/null'.format(shlex.quote(val.path)),
-                    'if [[ -e {0} ]]; then rm {0}; fi'.format(dest.remotepath),
-                    'mkfifo {}'.format(dest.remotepath),
+                    'if [[ -e {0} ]]; then rm {0}; fi'.format(dest),
+                    'mkfifo {}'.format(dest),
                     "gsutil {} cat {} > {} &".format(
                         '-u {}'.format(shlex.quote(self.project)) if self.get_requester_pays(val.path) else '',
                         shlex.quote(val.path),
-                        dest.remotepath
+                        dest
                     )
                 ]
                 exports.append('export {}="{}"'.format(
                     key,
-                    dest.remotepath
+                    dest
                 ))
             elif val.type == 'download':
                 job_vars.append(shlex.quote(key))
@@ -575,7 +581,7 @@ def job_setup_teardown(self, jobId: str, patterns: typing.Dict[str, str]) -> typ
                 'mkdir -p $CANINE_JOB_INPUTS',
                 'mkdir -p $CANINE_JOB_ROOT',
             ] + exports
-        ) + '\ncd $CANINE_JOB_ROOT\n'
+        ) + '\nexport CANINE_DOCKER_ARGS="{docker}"\ncd $CANINE_JOB_ROOT\n'.format(docker=' '.join(docker_args))
 
         # generate localization script
         localization_script = '\n'.join([
@@ -599,6 +605,7 @@ def job_setup_teardown(self, jobId: str, patterns: typing.Dict[str, str]) -> typ
                         for name, pattern in patterns.items()
                     )
                 ),
+                'if [[ -n "$CANINE_STREAM_DIR" ]]; then rm -rf $CANINE_STREAM_DIR; fi'
             ]
         )
         return setup_script, localization_script, teardown_script

diff --git a/canine/test/test_localizer_batched.py b/canine/test/test_localizer_batched.py
@@ -305,7 +305,7 @@ def test_setup_teardown(self):
                                 path = value.path
                                 if value.type == 'stream':
                                     src = path
-                                    path = localizer.reserve_path('jobs', str(jid), 'inputs', os.path.basename(os.path.abspath(src))).remotepath
+                                    path = os.path.join('$CANINE_STREAM_DIR', os.path.basename(os.path.abspath(src)))
                                     self.assertIn(
                                         'if [[ -e {dest} ]]; then rm {dest}; fi\n'
                                         'mkfifo {dest}\n'
@@ -329,6 +329,8 @@ def test_setup_teardown(self):
                                     )
                                 if isinstance(path, PathType):
                                     path = path.remotepath
+                                if '$' in path:
+                                    path = path.replace('$', '\\$')
                                 self.assertRegex(
                                     setup_text,
                                     r'export {}=[\'"]?{}[\'"]?'.format(arg, path)

diff --git a/examples/gpu.yaml b/examples/gpu.yaml
@@ -1,7 +1,7 @@
 name: GPU Test and example
 script:
   - nvidia-smi
-  - sudo docker run --rm --runtime nvidia nvidia/cuda nvidia-smi
+  - sudo docker run --rm --runtime nvidia $CANINE_DOCKER_ARGS nvidia/cuda nvidia-smi
 inputs:
   dummy_input: h
 resources:

diff --git a/examples/rnaseqc.yaml b/examples/rnaseqc.yaml
@@ -1,6 +1,6 @@
 name: RNA-SeQC
 script:
-  - sudo docker run --rm -v $CANINE_ROOT:$CANINE_ROOT -v $CANINE_JOB_ROOT:/output gcr.io/broad-cga-aarong-gtex/rnaseqc rnaseqc $genes_gtf $bam_file output --coverage -vv
+  - sudo docker run --rm $CANINE_DOCKER_ARGS gcr.io/broad-cga-aarong-gtex/rnaseqc rnaseqc $genes_gtf $bam_file $CANINE_JOB_ROOT --coverage -vv
 resources:
   cpus-per-task: 1
   mem-per-cpu: 3072M