add max-runtime argument

floydhub · Apr 16, 2018 · d127b0b · d127b0b
1 parent baed4ce
commit d127b0b
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/floyd/cli/run.py b/floyd/cli/run.py
@@ -156,9 +156,10 @@ def show_new_job_info(expt_client, job_name, expt_info, mode, open_notebook=True
 @click.option('--cpu+', 'cpup', is_flag=True, help='Run in a CPU+ instance')
 @click.option('--gpu2', 'gpu2', is_flag=True, help='Run in a GPU2 instance')
 @click.option('--cpu2', 'cpu2', is_flag=True, help='Run in a CPU2 instance')
+@click.option('--max-runtime', '-m', help='Max runtime to override for the job, in seconds')
 @click.argument('command', nargs=-1)
 @click.pass_context
-def run(ctx, gpu, env, message, data, mode, open_notebook, follow, tensorboard, gpup, cpup, gpu2, cpu2, command):
+def run(ctx, gpu, env, message, data, mode, open_notebook, follow, tensorboard, gpup, cpup, gpu2, cpu2, max_runtime, command):
     """
     Run a command on Floyd. Floyd will upload contents of the
     current directory and run your command remotely.
@@ -238,11 +239,14 @@ def run(ctx, gpu, env, message, data, mode, open_notebook, follow, tensorboard,
 
     # Create experiment request
     # Get the actual command entered in the command line
+    if max_runtime:
+        max_runtime = int(max_runtime)
     full_command = get_command_line(instance_type, env, message, data, mode, open_notebook, tensorboard, command_str)
     experiment_request = ExperimentRequest(name=experiment_name,
                                            description=message,
                                            full_command=full_command,
                                            module_id=module_id,
+                                           max_runtime=max_runtime,
                                            env=env,
                                            data_ids=data_ids,
                                            family_id=experiment_config.family_id,

diff --git a/floyd/model/experiment.py b/floyd/model/experiment.py
@@ -105,6 +105,7 @@ class ExperimentRequestSchema(Schema):
     family_id = fields.Str(allow_none=True)
     instance_type = fields.Str(allow_none=True)
     env = fields.Str(allow_none=True)
+    max_runtime = fields.Integer(allow_none=True)
 
     @post_load
     def make_experiment_request(self, kwargs):
@@ -122,7 +123,8 @@ def __init__(self,
                  env=None,
                  data_ids=None,
                  family_id=None,
-                 instance_type=None):
+                 instance_type=None,
+                 max_runtime=None):
         self.name = name
         self.description = description
         self.full_command = full_command
@@ -131,3 +133,4 @@ def __init__(self,
         self.family_id = family_id
         self.instance_type = instance_type
         self.env = env
+        self.max_runtime = max_runtime