2 files changed, 49 insertions, 23 deletions

diff --git a/app/relay/modules/samplernn.js b/app/relay/modules/samplernn.js
index 3cef25d..8b3b8e6 100644
--- a/app/relay/modules/samplernn.js
+++ b/app/relay/modules/samplernn.js
@@ -16,35 +16,29 @@ const train = {
   script: 'train.py',
   params: (task) => {
     return [
-      '--exp', task.checkpoint,
+      '--exp', task.dataset,
       '--dataset', task.dataset,
       '--frame_sizes', '8', '2',
       '--n_rnn', '2',
-      '--epoch_limit', task.opt.epoch_limit,
-      '--sample_length', task.opt.sample_length,
-      '--n_samples', task.opt.n_samples,
+      '--epoch_limit', task.epochs || 4,
+      '--sample_length', task.opt.sample_length || 44100 * 5,
+      '--n_samples', task.opt.n_samples || 6,
       '--keep_old_checkpoints', task.opt.keep_old_checkpoints ? 'True' : 'False',
     ]
   },
   onComplete: publish,
 }
-const report = {
-  type: 'perl',
-  script: 'latest.pl',
-  params: ['-v'],
-  isScript: true,
-}
 const generate = {
   type: 'pytorch',
   script: 'generate.py',
   params: (task) => {
     return [
-      '--exp', task.checkpoint,
+      '--exp', task.dataset,
       '--dataset', task.dataset,
       '--frame_sizes', '8', '2',
       '--n_rnn', '2',
-      '--sample_length', task.opt.sample_length,
-      '--n_samples', task.opt.n_samples,
+      '--sample_length', task.opt.sample_length || 44100 * 5,
+      '--n_samples', task.opt.n_samples || 6,
       '--keep_old_checkpoints', task.opt.keep_old_checkpoints ? 'True' : 'False',
     ]
   },
@@ -55,9 +49,14 @@ const publish = {
   script: 'latest.pl',
   params: (task) => {
     return ['-l', task.dataset]
-  } 
+  }
+}
+const report = {
+  type: 'perl',
+  script: 'latest.pl',
+  params: ['-v'],
+  isScript: true,
 }
-// after train and generate, run perl latest.pl -l $checkpoint_name
 
 export default {
   name, cwd,
diff --git a/app/relay/runner.js b/app/relay/runner.js
index deb3aa0..741ef8a 100644
--- a/app/relay/runner.js
+++ b/app/relay/runner.js
@@ -63,8 +63,7 @@ export function status () {
   }
 }
 
-export function build_params(module, task) {
-	const activity = module.activities[task.activity]
+export function build_params(module, activity, task) {
 	const interpreter = interpreters[activity.type]
   let opt_params;
 	if (typeof activity.params === 'function') {
@@ -159,7 +158,12 @@ export function run_script(task, cb) {
 export function run_task(task, preempt, watch){
 	const module = modules[task.module]
 	if (! module) return { type: 'error', error: "No such module: " + task.module }
-	const { activity, interpreter, params } = build_params(module, task)
+  const activity = module.activities[task.activity]
+  run_task_with_activity(task, module, activity, preempt, watch)
+}
+
+export function run_task_with_activity(task, module, activity, preempt, watch) {
+	const { interpreter, params } = build_params(module, task)
   if (! interpreter) return { type: 'error', error: "No such interpreter: " + activity.interpreter }
 
   if (interpreter.gpu && state.current_gpu_task.status !== 'IDLE') {
@@ -218,19 +222,42 @@ export function run_task(task, preempt, watch){
     if (finished) return
     finished = true
 		console.log('task error', subprocess.exitCode, err)
-    clear_task(interpreter.gpu, task)
-    remote.emit('task_res', { type: 'task_error', task, err })
-    set_connected(false)
+    finish({
+      type: 'task_error',
+      task, err,
+    })
 	})
 
 	subprocess.on('close', () => {
     if (finished) return
     finished = true
 		console.log('task ended', subprocess.exitCode || '')
-    clear_task(interpreter.gpu, task)
-    remote.emit('task_res', { type: 'task_finish', task })
     set_connected(false)
+    finish({
+      type: 'task_finish',
+      task,
+    })
 	})
+
+  function finish(task_res){
+    remote.emit('task_res', task_res)
+    clear_task(interpreter.gpu, task)
+    set_connected(false)
+    // remove task from queue
+    // queue.remove_task(task)
+    if (activity.onComplete) {
+      return run_task_with_activity(task, module, activity.onComplete, preempt, watch)
+    }
+    return run_next_task()
+  }
+}
+
+export function get_next_task(){
+  // get next task from the queue...
+  /*
+  const task = queue.get_next_task()
+  return run_task(task)
+  */
 }
 
 export function stop_task(task){