flask ping code updated

.github/workflows/deploy.yml

@@ -36,18 +36,18 @@ jobs:
         id: login-ecr
         uses: aws-actions/amazon-ecr-login@v1
 
-      - name: Build and push Docker image
-        env:
-          ECR_REGISTRY: ${{ steps.login-ecr.outputs.registry }}
-          ECR_REPOSITORY: customer-support-chatbot
-          IMAGE_TAG: latest
-        run: |
-          docker build \
-            --build-arg AWS_ACCESS_KEY_ID=${{ secrets.AWS_ACCESS_KEY_ID }} \
-            --build-arg AWS_SECRET_ACCESS_KEY=${{ secrets.AWS_SECRET_ACCESS_KEY }} \
-            --build-arg AWS_DEFAULT_REGION=${{ secrets.AWS_REGION }} \
-            -t $ECR_REGISTRY/$ECR_REPOSITORY:$IMAGE_TAG .
-          docker push $ECR_REGISTRY/$ECR_REPOSITORY:$IMAGE_TAG
+      # - name: Build and push Docker image
+      #   env:
+      #     ECR_REGISTRY: ${{ steps.login-ecr.outputs.registry }}
+      #     ECR_REPOSITORY: customer-support-chatbot
+      #     IMAGE_TAG: latest
+      #   run: |
+      #     docker build \
+      #       --build-arg AWS_ACCESS_KEY_ID=${{ secrets.AWS_ACCESS_KEY_ID }} \
+      #       --build-arg AWS_SECRET_ACCESS_KEY=${{ secrets.AWS_SECRET_ACCESS_KEY }} \
+      #       --build-arg AWS_DEFAULT_REGION=${{ secrets.AWS_REGION }} \
+      #       -t $ECR_REGISTRY/$ECR_REPOSITORY:$IMAGE_TAG .
+      #     docker push $ECR_REGISTRY/$ECR_REPOSITORY:$IMAGE_TAG
 
       - name: Deploy model to SageMaker
         run: |

app.py

@@ -169,6 +169,22 @@ def create_chat_interface():
 
     return interface
 
+from flask import Flask, jsonify
+import threading
+
+app = Flask(__name__)
+
+# Health check endpoint for SageMaker
+@app.route("/ping", methods=["GET"])
+def ping():
+    return jsonify({"status": "ok"}), 200
+
+# Start Flask app on a separate thread
+def run_flask():
+    app.run(host="0.0.0.0", port=8081)  # Different port for Flask
+
+threading.Thread(target=run_flask).start()
+
 if __name__ == "__main__":
     demo = create_chat_interface()
     print("Starting Gradient server...")

src/deploy_sagemaker.py

@@ -95,7 +95,11 @@ def deploy_app(acc_id, region_name, role_arn, ecr_repo_name, endpoint_name="cust
         logger.info(f"Deploying model to endpoint: {endpoint_name}")
         logger.info(f"Deployment configuration: {deployment_config}")
         
-        predictor = model.deploy(**deployment_config)
+        predictor = model.deploy(**deployment_config,health_check_timeout_seconds=180,  # Optionally, increase the timeout
+        environment={
+            "SAGEMAKER_CONTAINER_PORT": "8080",
+            "FLASK_HEALTHCHECK_PORT": "8081",
+        })
         
         logger.info(f"Successfully deployed to endpoint: {endpoint_name}")
         return predictor