fix: Change the library used in pytorch test causing cloudpickle version conflict (aws#1287)

Co-authored-by: Rohan Gujarathi <[email protected]>

Author: 2 people

tests/data/dummy_code_bundle_with_reqs/main_script.py

@@ -6,6 +6,6 @@
 import local_module
 
 print("Trying to import module from requirements.txt...")
-import stepfunctions
+from aws_xray_sdk.core import xray_recorder
 
 print("Done")

tests/data/dummy_code_bundle_with_reqs/requirements.txt

@@ -1,3 +1,3 @@
-# As a test dependency, we'll use the AWS Step Functions Data Science SDK - since it's more of a
-# notebook environment tool than a training/processing job library, so shouldn't be in base images
-stepfunctions
+# We use aws-xray-sdk as a test dependency as it shouldn't be in base images.
+# It is compatible with python 3.7 to 3.11, maintained regularly and does not depend on sagemaker
+aws-xray-sdk