Merge pull request #103 from nginyc/fix/stop_jobs_in_teardown

nudles · web-flow · commit 6744d1db984e · 2019-04-05T13:11:55.000+08:00
Stop jobs in teardown script
diff --git a/.env.sh b/.env.sh
@@ -10,6 +10,7 @@ export POSTGRES_EXT_PORT=5433
 export REDIS_EXT_PORT=6380
 export DATA_WORKDIR_PATH=$PWD/data # Shares a data folder with containers
 export LOGS_WORKDIR_PATH=$PWD/logs # Shares a folder with containers that stores components' logs
+export APP_MODE=DEV # DEV or PROD
 
 # Internal credentials for Rafiki's components
 export POSTGRES_USER=rafiki
diff --git a/examples/models/image_classification/SkDt.py b/examples/models/image_classification/SkDt.py
@@ -5,7 +5,6 @@
 import base64
 import numpy as np
 
-from rafiki.config import APP_MODE
 from rafiki.model import BaseModel, InvalidModelParamsException, test_model_class, \
                         IntegerKnob, CategoricalKnob, dataset_utils, logger
 from rafiki.constants import TaskType, ModelDependency
@@ -17,7 +16,7 @@ class SkDt(BaseModel):
     @staticmethod
     def get_knob_config():
         return {
-            'max_depth': IntegerKnob(2, 16 if APP_MODE != 'DEV' else 4),
+            'max_depth': IntegerKnob(2, 4),
             'criterion': CategoricalKnob(['gini', 'entropy'])
         }
 
diff --git a/examples/models/image_classification/SkSvm.py b/examples/models/image_classification/SkSvm.py
@@ -5,7 +5,6 @@
 import base64
 import numpy as np
 
-from rafiki.config import APP_MODE
 from rafiki.model import BaseModel, InvalidModelParamsException, test_model_class, \
                         IntegerKnob, CategoricalKnob, FloatKnob, dataset_utils
 from rafiki.constants import TaskType, ModelDependency
@@ -17,7 +16,7 @@ class SkSvm(BaseModel):
     @staticmethod
     def get_knob_config():
         return {
-            'max_iter': IntegerKnob(10, 40 if APP_MODE != 'DEV' else 10),
+            'max_iter': IntegerKnob(10, 20),
             'kernel': CategoricalKnob(['rbf', 'linear']),
             'gamma': CategoricalKnob(['scale', 'auto']),
             'C': FloatKnob(1e-2, 1e2, is_exp=True)
diff --git a/examples/models/image_classification/TfFeedForward.py b/examples/models/image_classification/TfFeedForward.py
@@ -7,7 +7,6 @@
 import numpy as np
 import base64
 
-from rafiki.config import APP_MODE
 from rafiki.model import BaseModel, InvalidModelParamsException, test_model_class, \
                         IntegerKnob, CategoricalKnob, FloatKnob, FixedKnob, dataset_utils, logger
 from rafiki.constants import TaskType, ModelDependency
@@ -20,8 +19,8 @@ class TfFeedForward(BaseModel):
     @staticmethod
     def get_knob_config():
         return {
-            'epochs': IntegerKnob(3, 10 if APP_MODE != 'DEV' else 3),
-            'hidden_layer_count': IntegerKnob(1, 8 if APP_MODE != 'DEV' else 2),
+            'epochs': FixedKnob(3),
+            'hidden_layer_count': IntegerKnob(1, 2),
             'hidden_layer_units': IntegerKnob(2, 128),
             'learning_rate': FloatKnob(1e-5, 1e-1, is_exp=True),
             'batch_size': CategoricalKnob([16, 32, 64, 128]),
diff --git a/examples/models/image_classification/TfVgg16.py b/examples/models/image_classification/TfVgg16.py
@@ -11,7 +11,6 @@
 from rafiki.model import BaseModel, InvalidModelParamsException, test_model_class, \
                         IntegerKnob, FloatKnob, CategoricalKnob, dataset_utils
 from rafiki.constants import TaskType, ModelDependency
-from rafiki.config import APP_MODE
 
 class TfVgg16(BaseModel):
     '''
@@ -20,7 +19,7 @@ class TfVgg16(BaseModel):
     @staticmethod
     def get_knob_config():
         return {
-            'epochs': IntegerKnob(1, 1 if APP_MODE != 'DEV' else 10),
+            'epochs': FixedKnob(1),
             'learning_rate': FloatKnob(1e-5, 1e-1, is_exp=True),
             'batch_size': CategoricalKnob([16, 32, 64, 128]),
         }
diff --git a/examples/models/pos_tagging/PyBiLstm.py b/examples/models/pos_tagging/PyBiLstm.py
@@ -15,7 +15,6 @@
 from rafiki.model import BaseModel, InvalidModelParamsException, test_model_class, \
                         IntegerKnob, FloatKnob, CategoricalKnob, logger, dataset_utils
 from rafiki.constants import TaskType, ModelDependency
-from rafiki.config import APP_MODE
 
 class PyBiLstm(BaseModel):
     '''
@@ -24,7 +23,7 @@ class PyBiLstm(BaseModel):
     @staticmethod
     def get_knob_config():
         return {
-            'epochs': IntegerKnob(10, 50 if APP_MODE != 'DEV' else 10),
+            'epochs': FixedKnob(10),
             'word_embed_dims': IntegerKnob(16, 128),
             'word_rnn_hidden_size': IntegerKnob(16, 128),
             'word_dropout': FloatKnob(1e-3, 2e-1, is_exp=True),
diff --git a/rafiki/admin/admin.py b/rafiki/admin/admin.py
@@ -6,8 +6,9 @@
 import csv
 
 from rafiki.db import Database
-from rafiki.constants import ServiceStatus, UserType, ServiceType, TrainJobStatus, ModelAccessRight, BudgetType
-from rafiki.config import MIN_SERVICE_PORT, MAX_SERVICE_PORT, SUPERADMIN_EMAIL, SUPERADMIN_PASSWORD
+from rafiki.constants import ServiceStatus, UserType, ServiceType, InferenceJobStatus, \
+    TrainJobStatus, ModelAccessRight, BudgetType
+from rafiki.config import SUPERADMIN_EMAIL, SUPERADMIN_PASSWORD
 from rafiki.model import ModelLogger
 from rafiki.container import DockerSwarmContainerManager 
 
@@ -288,6 +289,18 @@ def stop_train_job_worker(self, service_id):
             'sub_train_job_id': worker.sub_train_job_id
         }
 
+    def stop_all_train_jobs(self):
+        train_jobs = self._db.get_train_jobs_by_status(TrainJobStatus.RUNNING)
+        for train_job in train_jobs:
+            self._services_manager.stop_train_services(train_job.id)
+
+        return [
+            {
+                'id': train_job.id
+            }
+            for train_job in train_jobs
+        ]
+
     ####################################
     # Trials
     ####################################
@@ -467,6 +480,18 @@ def get_inference_jobs_by_user(self, user_id):
             for (inference_job, train_job, predictor_host) in zip(inference_jobs, train_jobs, predictor_hosts)
         ]
 
+    def stop_all_inference_jobs(self):
+        inference_jobs = self._db.get_inference_jobs_by_status(InferenceJobStatus.RUNNING)
+        for inference_job in inference_jobs:
+            self._services_manager.stop_inference_services(inference_job.id)
+            
+        return [
+            {
+                'id': inference_job.id
+            }
+            for inference_job in inference_jobs
+        ]
+
     ####################################
     # Models
     ####################################
diff --git a/rafiki/admin/app.py b/rafiki/admin/app.py
@@ -298,6 +298,23 @@ def get_models(auth):
         with admin:
             return jsonify(admin.get_models(auth['user_id'], **params))
 
+####################################
+# Administrative Actions
+####################################
+
+@app.route('/actions/stop_all_jobs', methods=['POST'])
+@auth([UserType.ADMIN])
+def stop_all_jobs(auth):
+    admin = get_admin()
+
+    with admin:
+        train_jobs = admin.stop_all_train_jobs()
+        inference_jobs = admin.stop_all_inference_jobs()
+        return jsonify({
+            'train_jobs': train_jobs,
+            'inference_jobs': inference_jobs
+        })
+    
 # Handle uncaught exceptions with a server error & the error's stack trace (for development)
 @app.errorhandler(Exception)
 def handle_error(error):
diff --git a/rafiki/admin/services_manager.py b/rafiki/admin/services_manager.py
@@ -2,11 +2,12 @@
 import logging
 import traceback
 import time
+import socket
+from contextlib import closing
 
 from rafiki.db import Database
 from rafiki.constants import ServiceStatus, UserType, ServiceType, BudgetType
-from rafiki.config import MIN_SERVICE_PORT, MAX_SERVICE_PORT, \
-    TRAIN_WORKER_REPLICAS_PER_SUB_TRAIN_JOB, INFERENCE_WORKER_REPLICAS_PER_TRIAL, \
+from rafiki.config import TRAIN_WORKER_REPLICAS_PER_SUB_TRAIN_JOB, INFERENCE_WORKER_REPLICAS_PER_TRIAL, \
     INFERENCE_MAX_BEST_TRIALS, SERVICE_STATUS_WAIT
 from rafiki.container import DockerSwarmContainerManager, ServiceRequirement, InvalidServiceRequest
 from rafiki.model import parse_model_install_command
@@ -122,9 +123,11 @@ def stop_train_services(self, train_job_id):
         train_job = self._db.get_train_job(train_job_id)
 
         # Stop all workers for train job
-        workers = self._db.get_workers_of_train_job(train_job_id)
-        for worker in workers:
-            self._stop_train_job_worker(worker)
+        sub_train_jobs = self._db.get_sub_train_jobs_of_train_job(train_job_id)
+        for sub_train_job in sub_train_jobs:
+            workers = self._db.get_workers_of_sub_train_job(sub_train_job.id)
+            for worker in workers:
+                self._stop_train_job_worker(worker)
 
         return train_job
         
@@ -345,19 +348,13 @@ def _create_service(self, service_type, docker_image,
 
         return service
 
-    # Compute next available external port
     def _get_available_ext_port(self):
-        services = self._db.get_services(status=ServiceStatus.RUNNING)
-        used_ports = [int(x.ext_port) for x in services if x.ext_port is not None]
-        port = MIN_SERVICE_PORT
-        while port <= MAX_SERVICE_PORT:
-            if port not in used_ports:
-                return port
-
-            port += 1
-
-        return port
-
+        # Credits to https://stackoverflow.com/questions/1365265/on-localhost-how-do-i-pick-a-free-port-number
+        with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
+            s.bind(('', 0))
+            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            return s.getsockname()[1]
+    
     def _get_best_trials_for_inference(self, inference_job):
         best_trials = self._db.get_best_trials_of_train_job(inference_job.train_job_id)
         return best_trials
diff --git a/rafiki/client/client.py b/rafiki/client/client.py
@@ -582,6 +582,19 @@ def delete_advisor(self, advisor_id):
         data = self._delete('/advisors/{}'.format(advisor_id), target='advisor')
         return data
 
+    ####################################
+    # Administrative Actions
+    ####################################
+
+    def stop_all_jobs(self):
+        '''
+        Stops all train and inference jobs on Rafiki. 
+
+        Only admins can call this.
+        '''
+        data = self._post('/actions/stop_all_jobs')
+        return data
+
     ####################################
     # Private
     ####################################
diff --git a/rafiki/config.py b/rafiki/config.py
@@ -2,13 +2,10 @@
 
 # Global
 APP_SECRET = 'rafiki'
-APP_MODE = 'DEV' # DEV or PROD
 SUPERADMIN_EMAIL = 'superadmin@rafiki'
 SUPERADMIN_PASSWORD = 'rafiki'
 
 # Admin
-MIN_SERVICE_PORT = 30000
-MAX_SERVICE_PORT = 32767
 SERVICE_STATUS_WAIT = 1
 TRAIN_WORKER_REPLICAS_PER_SUB_TRAIN_JOB = 2
 INFERENCE_WORKER_REPLICAS_PER_TRIAL = 2
diff --git a/rafiki/db/database.py b/rafiki/db/database.py
@@ -1,6 +1,6 @@
 import datetime
 import os
-from sqlalchemy import create_engine
+from sqlalchemy import create_engine, distinct
 from sqlalchemy.orm import sessionmaker
 
 from rafiki.constants import TrainJobStatus, \
@@ -84,6 +84,12 @@ def get_train_job(self, id):
         train_job = self._session.query(TrainJob).get(id)
         return train_job
 
+    def get_train_jobs_by_status(self, status):
+        job_ids = self._session.query(distinct(SubTrainJob.train_job_id)) \
+            .filter(SubTrainJob.status == status).all()
+        return self._session.query(TrainJob) \
+            .filter(TrainJob.id.in_(job_ids)).all()
+
     # Returns for the latest app version unless specified
     def get_train_job_by_app_version(self, app, app_version=-1):
         # pylint: disable=E1111
@@ -217,6 +223,12 @@ def get_inference_jobs_of_app(self, app):
             .order_by(InferenceJob.datetime_started.desc()).all()
 
         return inference_jobs
+    
+    def get_inference_jobs_by_status(self, status):
+        jobs = self._session.query(InferenceJob) \
+            .filter(InferenceJob.status == status).all()
+
+        return jobs
 
     ####################################
     # Inference Job Workers
diff --git a/scripts/start_admin.sh b/scripts/start_admin.sh
@@ -1,3 +1,10 @@
+# Mount whole project folder with code for dev for shorter iterations
+if [ $APP_MODE = "DEV" ]; then
+  VOLUME_MOUNTS="-v $PWD:$DOCKER_WORKDIR_PATH"
+else
+  VOLUME_MOUNTS="-v $LOGS_WORKDIR_PATH:$LOGS_DOCKER_WORKDIR_PATH -v $DATA_WORKDIR_PATH:$DATA_DOCKER_WORKDIR_PATH"
+fi
+
 docker run --rm --name $ADMIN_HOST \
   --network $DOCKER_NETWORK \
   -e POSTGRES_HOST=$POSTGRES_HOST \
@@ -22,7 +29,6 @@ docker run --rm --name $ADMIN_HOST \
   -e DATA_DOCKER_WORKDIR_PATH=$DATA_DOCKER_WORKDIR_PATH \
   -e DOCKER_WORKDIR_PATH=$DOCKER_WORKDIR_PATH \
   -v /var/run/docker.sock:/var/run/docker.sock \
-  -v $LOGS_WORKDIR_PATH:$LOGS_DOCKER_WORKDIR_PATH \
-  -v $DATA_WORKDIR_PATH:$DATA_DOCKER_WORKDIR_PATH \
+  $VOLUME_MOUNTS \
   -p $ADMIN_EXT_PORT:$ADMIN_PORT \
   $RAFIKI_IMAGE_ADMIN:$RAFIKI_VERSION
diff --git a/scripts/start_advisor.sh b/scripts/start_advisor.sh
@@ -1,10 +1,16 @@
+# Mount whole project folder with code for dev for shorter iterations
+if [ $APP_MODE = "DEV" ]; then
+  VOLUME_MOUNTS="-v $PWD:$DOCKER_WORKDIR_PATH"
+else
+  VOLUME_MOUNTS="-v $LOGS_WORKDIR_PATH:$LOGS_DOCKER_WORKDIR_PATH -v $DATA_WORKDIR_PATH:$DATA_DOCKER_WORKDIR_PATH"
+fi
+
 docker run --rm --name $ADVISOR_HOST \
   --network $DOCKER_NETWORK \
   -e LOGS_WORKDIR_PATH=$LOGS_WORKDIR_PATH \
   -e DATA_WORKDIR_PATH=$DATA_WORKDIR_PATH \
   -e LOGS_DOCKER_WORKDIR_PATH=$LOGS_DOCKER_WORKDIR_PATH \
   -e DATA_DOCKER_WORKDIR_PATH=$DATA_DOCKER_WORKDIR_PATH \
-  -v $LOGS_WORKDIR_PATH:$LOGS_DOCKER_WORKDIR_PATH \
-  -v $DATA_WORKDIR_PATH:$DATA_DOCKER_WORKDIR_PATH \
+  $VOLUME_MOUNTS \
   -p $ADVISOR_EXT_PORT:$ADVISOR_PORT \
   $RAFIKI_IMAGE_ADVISOR:$RAFIKI_VERSION
diff --git a/scripts/stop.sh b/scripts/stop.sh
@@ -12,6 +12,9 @@ title()
     echo "$edge"
 }
 
+title "Stopping any existing jobs..."
+python scripts/stop_all_jobs.py
+
 title "Stopping Rafiki's DB..."
 docker rm -f $POSTGRES_HOST || echo "Failed to stop Rafiki's DB"
 
diff --git a/scripts/stop_all_jobs.py b/scripts/stop_all_jobs.py
@@ -0,0 +1,15 @@
+import os
+
+from rafiki.client import Client
+from rafiki.config import SUPERADMIN_EMAIL, SUPERADMIN_PASSWORD
+
+if __name__ == '__main__':
+    rafiki_host = os.environ.get('RAFIKI_HOST', 'localhost')
+    admin_port = int(os.environ.get('ADMIN_EXT_PORT', 3000))
+    user_email = os.environ.get('USER_EMAIL', SUPERADMIN_EMAIL)
+    user_password = os.environ.get('USER_PASSWORD', SUPERADMIN_PASSWORD)
+
+    # Initialize client
+    client = Client(admin_host=rafiki_host, admin_port=admin_port)
+    client.login(email=user_email, password=user_password)
+    print(client.stop_all_jobs())
diff --git a/web/package.json b/web/package.json
@@ -6,8 +6,9 @@
   "author": "nginyc",
   "license": "Apache-2.0",
   "scripts": {
-    "build": "webpack --config webpack.config.js",
-    "dev": "webpack-dev-server --open --config webpack.config.js"
+    "build": "webpack --config webpack.config.prod.js",
+    "start": "node app.js",
+    "dev": "webpack-dev-server --open --config webpack.config.dev.js"
   },
   "dependencies": {
     "@material-ui/core": "^3.3.2",
diff --git a/web/webpack.config.dev.js b/web/webpack.config.dev.js
@@ -21,8 +21,8 @@ module.exports = {
   plugins: [
     new CleanWebpackPlugin(['dist']),
     new webpack.DefinePlugin({
-      'process.env.NODE_ENV': JSON.stringify('development'),
-      'process.env.VERSION': JSON.stringify(packageJson.version)
+      'window.ADMIN_HOST': JSON.stringify(process.env.RAFIKI_ADDR),
+      'window.ADMIN_PORT': JSON.stringify(process.env.ADMIN_EXT_PORT)
     })
   ],
   module: {
diff --git a/web/webpack.config.prod.js b/web/webpack.config.prod.js
