Add node availability zone variable for aws (#25)

* add node availability zone variable for aws
Signed-off-by: vsoch <vsoch@users.noreply.github.com>

Author: vsoch

docs/getting_started/aws.md

@@ -66,6 +66,10 @@ variables:
 
     # Customize availability zones for this experiment
     availability_zones: [us-east-1a, us-east-1b]
+
+    # Important for instance types only in one zone (hpc instances)
+    # Select your node group availability zone:
+    node_group_availability_zone: us-east-2b
 ```
 
 Note that we currently take a simple approach for boolean values - if it's present (e.g., the examples)

fluxcloud/client/config.py

@@ -12,7 +12,6 @@
 
 
 def main(args, parser, extra, subparser):
-
     fluxcloud.utils.ensure_no_extra(extra)
 
     # If nothing provided, show help

fluxcloud/client/shell.py

@@ -8,7 +8,6 @@
 
 
 def main(args, parser, extra, subparser):
-
     fluxcloud.utils.ensure_no_extra(extra)
     lookup = {"ipython": ipython, "python": python, "bpython": bpython}
     shells = ["ipython", "python", "bpython"]
@@ -32,7 +31,6 @@ def main(args, parser, extra, subparser):
 
 
 def create_client(args):
-
     cli = Client(
         quiet=args.quiet,
         settings_file=args.settings_file,

fluxcloud/logger.py

@@ -38,7 +38,6 @@ def add_prefix(msg, char=">>"):
 
 
 class ColorizingStreamHandler(_logging.StreamHandler):
-
     BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)
     RESET_SEQ = LogColors.ENDC
     COLOR_SEQ = "\033[%dm"

fluxcloud/main/api.py

@@ -84,7 +84,6 @@ def submit(self, setup, experiment, size):
         # Sleep time will be time of last job, assuming they are similar
         sleep_time = 5
         for jobname, job in experiment.jobs.items():
-
             # Do we want to run this job for this size and machine?
             if not experiment.check_job_run(job, size):
                 logger.debug(

fluxcloud/main/client.py

@@ -181,7 +181,6 @@ def submit(self, setup, experiment):
 
         # Iterate through all the cluster sizes
         for size in experiment.minicluster["size"]:
-
             # We can't run if the minicluster > the experiment size
             if size > experiment.size:
                 logger.warning(
@@ -223,7 +222,6 @@ def apply(self, setup, experiment):
         experiment_dir = experiment.root_dir
 
         for size, jobname, job in experiment.iter_jobs():
-
             # Add the size
             jobname = f"{jobname}-minicluster-size-{size}"
             job_output = os.path.join(experiment_dir, jobname)
@@ -265,13 +263,15 @@ def apply(self, setup, experiment):
             # Apply the job, and save to output directory
             self.run_timed(f"{self.job_prefix}-{jobname}", ["/bin/bash", apply_script])
 
+            # Save times between experiment runs
+            experiment.save_metadata(self.times, self.info)
+
     def clear_minicluster_times(self):
         """
         Update times to not include jobs
         """
         times = {}
         for key, value in self.times.items():
-
             # Don't add back a job that was already saved
             if key.startswith(self.job_prefix):
                 continue

fluxcloud/main/clouds/aws/templates/cluster-config.yaml

@@ -16,6 +16,7 @@ managedNodeGroups:
     minSize: {{ experiment.size }}
     maxSize: {{ experiment.size }}
     labels: { "fluxoperator": "true" }
+    {% if variables['node_group_availability_zone'] %}availabilityZones: ["{{ variables['node_group_availability_zone'] }}"]{% endif %}
     {% if variables['ssh_key'] %}ssh:
       allow: true
       publicKeyPath: {{ variables['ssh_key'] }}{% endif %}

fluxcloud/main/decorator.py

@@ -22,7 +22,6 @@ class save_meta(Decorator):
     """
 
     def __call__(self, cls, *args, **kwargs):
-
         # Name of the key is after command
         idx = 0
         if "setup" in kwargs:
@@ -52,7 +51,6 @@ class timed(Decorator):
     """
 
     def __call__(self, cls, *args, **kwargs):
-
         # Name of the key is after command
         if "name" in kwargs:
             key = kwargs["name"]

fluxcloud/main/experiment.py

@@ -173,7 +173,6 @@ def iter_jobs(self):
 
         # Iterate through all the cluster sizes
         for size in minicluster["size"]:
-
             # We can't run if the minicluster > the experiment size
             if size > self.size:
                 logger.warning(
@@ -183,7 +182,6 @@ def iter_jobs(self):
 
             # Jobname is used for output
             for jobname, job in self.jobs.items():
-
                 # Do we want to run this job for this size and machine?
                 if not self.check_job_run(job, size):
                     logger.debug(
@@ -199,7 +197,6 @@ def get_persistent_variables(self, size, required=None):
         """
         jobvars = {}
         for _, job in self.jobs.items():
-
             # Skip jobs targeted for a different size
             if "size" in job and job["size"] != size:
                 continue
@@ -302,7 +299,6 @@ def is_run(self):
 
         # If all job output files exist, experiment is considered run
         for size in self.minicluster["size"]:
-
             # We can't run if the minicluster > the experiment size
             if size > self.size:
                 logger.warning(
@@ -312,7 +308,6 @@ def is_run(self):
 
             # Jobname is used for output
             for jobname, job in self.jobs.items():
-
                 # Do we want to run this job for this size and machine?
                 if not self.check_job_run(job, size):
                     logger.debug(

setup.py

@@ -65,7 +65,6 @@ def get_reqs(lookup=None, key="INSTALL_REQUIRES"):
 
 
 if __name__ == "__main__":
-
     INSTALL_REQUIRES = get_reqs(lookup)
     TESTS_REQUIRES = get_reqs(lookup, "TESTS_REQUIRES")
     INSTALL_REQUIRES_ALL = get_reqs(lookup, "INSTALL_REQUIRES_ALL")