Merge pull request #167 from tgen/improve_gpu_detection_and_update_se…

…tup_py Small update to slurm_singularity's auto gpu detection and configuring setup.py to use networkx>=2.6,<3
tgen · Nov 17, 2023 · ecd8dfd · ecd8dfd
2 parents 93ef04d + 57dc5ff
commit ecd8dfd
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/jetstream/backends/slurm_singularity.py b/jetstream/backends/slurm_singularity.py
@@ -673,8 +673,8 @@ async def sbatch(cmd, identity, singularity_image, singularity_executable="singu
         singularity_exec_args += " --bind {}".format(os.getenv('JS_PIPELINE_PATH'))
         sbatch_script += "export SINGULARITYENV_JS_PIPELINE_PATH={}\n".format(os.getenv('JS_PIPELINE_PATH'))
 
-    if any('gpu' in s for s in [singularity_args, sbatch_args]):
-        if all('--nv' not in s for s in singularity_args):
+    if 'gpu' in ' '.join(map(str, sbatch_args)):
+        if '--nv' not in singularity_exec_args:
             singularity_exec_args += ' --nv'
 
     for arg in singularity_args:

diff --git a/setup.py b/setup.py
@@ -38,7 +38,7 @@
         'Topic :: Utilities',
     ],
     install_requires=[
-        'networkx',
+        'networkx>=2.6,<3',
         'pyyaml',
         'ulid-py',
         'jinja2>=3.0.0',