dep fixes

Author: Bryannah Hernandez

requirements/extras/feature-processor_requirements.txt

@@ -1,2 +1,6 @@
 pyspark==3.3.1
 sagemaker-feature-store-pyspark-3.3
+uvicorn>=0.30.1
+fastapi>=0.111.0
+nest-asyncio
+transformers

requirements/extras/huggingface_requirements.txt

@@ -3,4 +3,4 @@ sagemaker_schema_inference_artifacts>=0.0.5
 uvicorn>=0.30.1
 fastapi>=0.111.0
 nest-asyncio
-transformers>=3.0.2
+transformers

requirements/extras/test_requirements.txt

@@ -39,3 +39,7 @@ schema==0.7.5
 tensorflow>=2.1,<=2.16
 mlflow>=2.12.2,<2.13
 huggingface_hub>=0.23.4
+uvicorn>=0.30.1
+fastapi>=0.111.0
+nest-asyncio
+transformers

src/sagemaker/serve/app.py

@@ -9,15 +9,6 @@
 from transformers import pipeline
 from fastapi import FastAPI, Request
 
-if not importlib.util.find_spec("uvicorn"):
-    raise ImportError("Unable to import uvicorn, check if uvicorn is installed")
-
-if not importlib.util.find_spec("transformers"):
-    raise ImportError("Unable to import transformers, check if transformers is installed")
-
-if not importlib.util.find_spec("fastapi"):
-    raise ImportError("Unable to import fastapi, check if fastapi is installed")
-
 
 logger = logging.getLogger(__name__)
 
@@ -58,6 +49,15 @@ def post(payload: dict):
 
 async def main():
     """Running server locally with uvicorn"""
+    if not importlib.util.find_spec("uvicorn"):
+        raise ImportError("Unable to import uvicorn, check if uvicorn is installed")
+
+    if not importlib.util.find_spec("transformers"):
+        raise ImportError("Unable to import transformers, check if transformers is installed")
+
+    if not importlib.util.find_spec("fastapi"):
+        raise ImportError("Unable to import fastapi, check if fastapi is installed")
+
     logger.info("Running")
     config = uvicorn.Config(
         "sagemaker.app:app",

tests/unit/sagemaker/serve/mode/test_in_process_mode.py

@@ -98,6 +98,12 @@ def test_load_ex(self, mock_session, mock_inference_spec, mock_path):
     def test_create_server_happy(
         self, mock_session, mock_inference_spec, mock_predictor, mock_logger
     ):
+        mock_start_serving = Mock()
+        mock_start_serving.side_effect = lambda *args, **kwargs: (
+            True,
+            None,
+        )
+
         mock_response = "Fake response"
         mock_multi_model_server_deep_ping = Mock()
         mock_multi_model_server_deep_ping.side_effect = lambda *args, **kwargs: (
@@ -114,6 +120,7 @@ def test_create_server_happy(
         )
 
         in_process_mode._multi_model_server_deep_ping = mock_multi_model_server_deep_ping
+        in_process_mode._start_serving = mock_start_serving
 
         in_process_mode.create_server(predictor=mock_predictor)
 
@@ -133,6 +140,12 @@ def test_create_server_ex(
         mock_inference_spec,
         mock_predictor,
     ):
+        mock_start_serving = Mock()
+        mock_start_serving.side_effect = lambda *args, **kwargs: (
+            True,
+            None,
+        )
+
         mock_multi_model_server_deep_ping = Mock()
         mock_multi_model_server_deep_ping.side_effect = lambda *args, **kwargs: (
             False,
@@ -148,5 +161,6 @@ def test_create_server_ex(
         )
 
         in_process_mode._multi_model_server_deep_ping = mock_multi_model_server_deep_ping
+        in_process_mode._start_serving = mock_start_serving
 
         self.assertRaises(LocalDeepPingException, in_process_mode.create_server, mock_predictor)