More env work

fwhigh · Jan 31, 2019 · 1705fd9 · 1705fd9
1 parent 826de26
commit 1705fd9
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/pmip/routes.py b/pmip/routes.py
@@ -28,7 +28,7 @@ def possible_types(value):
     )
     model = load_from_s3_and_unpickle(
         filename='model.pkl',
-        subdirectory=f'models/{latest_model_id}',
+        subdirectory=f'models/{os.getenv("ENVIRONMENT")}/{latest_model_id}',
         bucket=os.getenv('BUCKET')
     )
 
@@ -99,10 +99,10 @@ def post(self):
 @api.route('/model-info')
 class ModelInfo(Resource):
 
-    # @api.marshal_with(request, code=201)
     def get(self):
         result = {
-            "model_id": latest_model_id
+            "model_id": latest_model_id,
+            "ENVIRONMENT": os.getenv('ENVIRONMENT', '')
         }
 
         return {'result': result}, 200

diff --git a/scripts/train.sh b/scripts/train.sh
@@ -3,7 +3,7 @@
 set -e
 
 MODEL_ID=`date +%Y%m%d`
-S3_DIR=s3://$BUCKET/models/$MODEL_ID
+S3_DIR=s3://$BUCKET/models/$ENVIRONMENT/$MODEL_ID
 DIR=data
 
 echo MODEL_ID $MODEL_ID
@@ -28,7 +28,7 @@ papermill notebooks/model-training.ipynb $DIR/model-training-$MODEL_ID.ipynb -p
 jupyter nbconvert --to html $DIR/model-training-$MODEL_ID.ipynb
 
 # Push any assets to the cloud
-if [ "$ENVIRONMENT" == "staging" ]; then
+if [ "$ENVIRONMENT" == "staging" ] || [ "$ENVIRONMENT" == "prod" ]; then
     echo Pushing model to S3
     aws s3 cp $DIR/ $S3_DIR/ \
          --recursive --exclude "*" --include "*.ipynb" --include "*.html" --include "*.pkl"