feat: install sagemaker-training and sagemaker-inference libs in the image instead of deprecated sagemaker-containers

requirements.txt

@@ -1,32 +1,32 @@
-Flask==1.1.1  # sagemaker-containers requires flask 1.1.1
+Flask==2.2.5  # sagemaker-containers requires flask 1.1.1
 PyYAML==5.4.1
-Pillow==9.1.1
-boto3==1.17.52
-botocore==1.20.52
-cryptography==39.0.1
+Pillow==10.2.0
+boto3==1.28.57
+botocore==1.31.57
+cryptography==42.0.4
 dask==2022.11.1
 dask-cuda==22.12.0
-gunicorn==19.10.0
+gunicorn==22.0.0
 itsdangerous==2.0.1
 matplotlib==3.4.1
 multi-model-server==1.1.2
 numpy==1.24.1
 pandas==1.2.4
-protobuf==3.20.1
+protobuf==3.20.2
 psutil==5.6.7  # sagemaker-containers requires psutil 5.6.7
 pynvml==11.4.1
 python-dateutil==2.8.1
 retrying==1.3.3
 requests==2.29.0
-sagemaker-containers==2.8.6.post2
-sagemaker-inference==1.5.5
-scikit-learn==0.24.1
+sagemaker-training==4.7.4
+sagemaker-inference==1.10.1
+scikit-learn==1.2.1
 scipy==1.8.0
 urllib3==1.26.5
-wheel==0.36.2
+wheel==0.38.1
 jinja2==2.11.3
 MarkupSafe==1.1.1
-Werkzeug==0.15.6
+Werkzeug==3.0.3
 certifi==2023.7.22
 gevent==23.9.1
 numba==0.58.1

src/sagemaker_xgboost_container/algorithm_mode/__init__.py

@@ -12,12 +12,12 @@
 # language governing permissions and limitations under the License.
 import os
 
-from sagemaker_containers.beta.framework import env
+from sagemaker_inference import environment
 
 from sagemaker_xgboost_container.algorithm_mode import serve
 
 # Pre-load the model in the algorithm mode.
 # Otherwise, the model will be loaded when serving the first request per worker.
 # When the model is large, the request may timeout.
-if os.environ.get("SERVER_SOFTWARE") is not None and env.ServingEnv().module_name is None:
+if os.environ.get("SERVER_SOFTWARE") is not None and environment.Environment().module_name is None:
     serve.load_model()

src/sagemaker_xgboost_container/data_utils.py

@@ -25,7 +25,7 @@
 from mlio.integ.arrow import as_arrow_file
 from mlio.integ.numpy import as_numpy
 from mlio.integ.scipy import to_coo_matrix
-from sagemaker_containers import _content_types
+from sagemaker_training import content_types
 from scipy.sparse import vstack as scipy_vstack
 
 from sagemaker_algorithm_toolkit import exceptions as exc
@@ -45,7 +45,7 @@
     LIBSVM,
     PARQUET,
     RECORDIO_PROTOBUF,
-    _content_types.CSV,
+    content_types.CSV,
     xgb_content_types.LIBSVM,
     xgb_content_types.X_LIBSVM,
     xgb_content_types.X_PARQUET,
@@ -56,7 +56,7 @@
     CSV,
     PARQUET,
     RECORDIO_PROTOBUF,
-    _content_types.CSV,
+    content_types.CSV,
     xgb_content_types.X_PARQUET,
     xgb_content_types.X_RECORDIO_PROTOBUF,
 ]
@@ -102,7 +102,7 @@ def get_content_type(content_type_cfg_val):
         # the tuple ('text/csv', {'label_size': '1', 'charset': 'utf8'})
         content_type, params = cgi.parse_header(content_type_cfg_val.lower())
 
-        if content_type in [CSV, _content_types.CSV]:
+        if content_type in [CSV, content_types.CSV]:
             # CSV content type allows a label_size parameter
             # that should be 1 for XGBoost
             if params and "label_size" in params and params["label_size"] != "1":

src/sagemaker_xgboost_container/encoder.py

@@ -26,7 +26,7 @@
 import xgboost as xgb
 from mlio.integ.numpy import as_numpy
 from mlio.integ.scipy import to_coo_matrix
-from sagemaker_containers import _content_types, _errors
+from sagemaker_training import content_types, errors
 from scipy.sparse import vstack as scipy_vstack
 
 from sagemaker_xgboost_container.constants import xgb_content_types
@@ -103,7 +103,7 @@ def recordio_protobuf_to_dmatrix(string_like):  # type: (bytes) -> xgb.DMatrix
 
 
 _dmatrix_decoders_map = {
-    _content_types.CSV: csv_to_dmatrix,
+    content_types.CSV: csv_to_dmatrix,
     xgb_content_types.LIBSVM: libsvm_to_dmatrix,
     xgb_content_types.X_LIBSVM: libsvm_to_dmatrix,
     xgb_content_types.X_RECORDIO_PROTOBUF: recordio_protobuf_to_dmatrix,
@@ -142,4 +142,4 @@ def decode(obj, content_type):
         decoder = _dmatrix_decoders_map[media_content_type]
         return decoder(obj)
     except KeyError:
-        raise _errors.UnsupportedFormatError(media_content_type)
+        raise errors.UnsupportedFormatError(media_content_type)

test/resources/versions/train.py

@@ -21,8 +21,8 @@
 pyarrow==14.0.1
 python-dateutil==2.8.1
 retrying==1.3.3
-sagemaker-containers==2.8.6.post2
-sagemaker-inference==1.5.5
+sagemaker-training==4.7.4
+sagemaker-inference==1.10.1
 scikit-learn==0.24.1
 scipy==1.8.0
 smdebug==1.0.29

test/unit/algorithm_mode/test_serve_utils.py

@@ -19,8 +19,8 @@
 import numpy as np
 import pytest
 from mock import MagicMock, patch
-from sagemaker_containers._recordio import _read_recordio
-from sagemaker_containers.record_pb2 import Record
+from sagemaker_training.recordio import _read_recordio
+from sagemaker_training.record_pb2 import Record
 
 from sagemaker_algorithm_toolkit import exceptions as exc
 from sagemaker_xgboost_container import data_utils

test/unit/test_handler_service.py

@@ -16,7 +16,7 @@
 import pytest
 import xgboost as xgb
 from mock import patch
-from sagemaker_containers.beta.framework import content_types, encoders, errors
+from sagemaker_training import content_types, encoders, errors
 
 from sagemaker_xgboost_container.handler_service import HandlerService
 

test/unit/test_serving.py

@@ -18,7 +18,7 @@
 import pytest
 import xgboost as xgb
 from mock import MagicMock, patch
-from sagemaker_containers.beta.framework import content_types, encoders, errors
+from sagemaker_inference import content_types, encoders, errors
 
 from sagemaker_algorithm_toolkit.exceptions import UserError
 from sagemaker_xgboost_container import serving

test/utils/local_mode.py

@@ -27,7 +27,7 @@
 import yaml
 from botocore.exceptions import ClientError
 from sagemaker import fw_utils, utils
-from sagemaker_containers.beta.framework import content_types
+from sagemaker_training import content_types
 
 CYAN_COLOR = "\033[36m"
 END_COLOR = "\033[0m"