Find flux python

ats/atsMachines/fluxScheduled.py

@@ -10,15 +10,14 @@
 """
 
 import os
+import sys
 import time
 from math import ceil
 
-import flux
-import flux.job
-
 from ats import terminal
 from ats.atsMachines import lcMachines
 from ats.tests import AtsTest
+from ats.util.generic_utils import runCommand
 from ats import configuration
 from ats import log
 
@@ -40,6 +39,11 @@ def init(self):
         Sets ceiling on number of nodes and cores in the allocation.
         Defines a persistent handle to use to connect to the broker.
         """
+        # Add the flux python library to the system path
+        flux_python_path = runCommand("flux config builtin python_path")
+        sys.path.append(flux_python_path[0].strip())
+        import flux
+        import flux.job
         self.fluxHandle = flux.Flux()
         self.numNodes = int(
             flux.resource.list.resource_list(self.fluxHandle).get().up.nnodes
@@ -512,7 +516,7 @@ def remainingCapacity(self):
             return 0
         elif self.numProcsAvailable < 1:
             return 0
-        elif self.numGPUsAvailable < 1 and self.numGPUs is not 0:
+        elif self.numGPUsAvailable < 1 and self.numGPUs != 0:
             return 0
         else:
             return self.numProcsAvailable

ats/configuration.py

@@ -454,18 +454,22 @@ def get_machine_factory(module_name, machine_class,
     module_name. "machine_package" tells Python where module_name can be found if
     not in the project's root directory.
     """
-    log("Machine Factory: importing {} from {}".format(module_name, machine_package),
+    log(f"Machine Factory: importing {module_name} from {machine_package}",
         echo=False)
     try:
         machine_module = importlib.import_module(f'.{module_name}',
                                                  package=machine_package)
         machine_factory = getattr(machine_module, machine_class)
         return machine_factory
 
-    except ModuleNotFoundError:
-        log(f"Module '{module_name}' not found in package '{machine_package}'. Continuing search.",
-            echo=False)
-        return None
+    except ModuleNotFoundError as e:
+        if (e == ModuleNotFoundError(module_name)):
+            log(f"Module {module_name} not found in package {machine_package}. Continuing search.",
+                echo=False)
+            return None
+        else:
+            # If a module error occurs for a module other than module_name, raise an error
+            raise e
 
 def get_machine(file_text, file_name, is_batch=False):
     header = '#BATS:' if is_batch else '#ATS:'
@@ -524,12 +528,7 @@ def get_machine_entry_points(machine_class):
         echo=False)
     for name, machine_factory in ats_machines.items():
         if machine_class in machine_factory.value:
-            log("Machine Factory: Found machine {} of class {}: {}".format(
-                name,
-                machine_class,
-                machine_factory
-            ))
-
+            log(f"Machine Factory: Found machine {name} of class {machine_class}: {machine_factory}")
             return machine_factory.load()(machine_class, -1)
 
     # Downstream needs to be able to detect if machine isn't found

ats/util/generic_utils.py

@@ -358,7 +358,7 @@ def set_machine_type_based_on_sys_type():
         elif host.startswith('herd'):
             os.environ['MACHINE_TYPE'] = 'slurm32'
 
-        elif host.startswith('rzgenie') or host.startswith('rztopaz') or host.startswith('rztrona') or \
+        elif host.startswith('rzgenie') or host.startswith('rztrona') or \
              host.startswith('borax') or host.startswith('quartz') or host.startswith('agate') or \
              host.startswith('pascal') or host.startswith('jade') or host.startswith('mica'):
             os.environ['MACHINE_TYPE'] = 'slurm36'
@@ -372,6 +372,11 @@ def set_machine_type_based_on_sys_type():
         elif host.startswith('mammoth'):
             os.environ['MACHINE_TYPE'] = 'slurm128'
 
+        elif host.startswith('tioga') or host.startswith('rzadams') or \
+             host.startswith('rzvernal') or host.startswith('tuolumne') or \
+             host.startswith('corona'):
+            os.environ['MACHINE_TYPE'] = 'flux00'
+
         elif os.environ['SYS_TYPE'] in ['bgqos_0']:
             os.environ['MACHINE_TYPE'] = 'bgqos_0_ASQ'