feature: Add sagemaker-mlflow package

build_artifacts/v2/v2.0/v2.0.0/cpu.additional_packages_env.in

@@ -1,2 +1,3 @@
 #Add any additional packages here
-conda-forge::mlflow[version='>=2.13.0,<3.0']
+conda-forge::mlflow[version='>=2.13.2,<3.0']
+conda-forge::sagemaker-mlflow[version='>=0.1.0,<1.0']

build_artifacts/v2/v2.0/v2.0.0/gpu.additional_packages_env.in

@@ -1,2 +1,3 @@
 #Add any additional packages here
 conda-forge::mlflow[version='>=2.13.2,<3.0']
+conda-forge::sagemaker-mlflow[version='>=0.1.0,<1.0']

test/test_artifacts/v2/sagemaker-mlflow.test.Dockerfile

@@ -0,0 +1,19 @@
+ARG SAGEMAKER_DISTRIBUTION_IMAGE
+FROM $SAGEMAKER_DISTRIBUTION_IMAGE
+
+ARG MAMBA_DOCKERFILE_ACTIVATE=1
+
+RUN python -c "import sagemaker_mlflow"
+
+RUN sudo apt-get update && sudo apt-get install -y git && \
+    git clone --recursive https://github.com/aws/sagemaker-mlflow.git && \
+    :
+
+# For running sagemaker-mlflow tests, we need pytest
+RUN micromamba install -y --freeze-installed -c conda-forge pytest
+
+WORKDIR "sagemaker-mlflow/"
+COPY --chown=$MAMBA_USER:$MAMBA_USER scripts/run_sagemaker_mlflow_tests.sh .
+RUN chmod +x run_sagemaker_mlflow_tests.sh
+# Run tests in run_matplotlib_tests.sh
+CMD ["./run_sagemaker_mlflow_tests.sh"]

test/test_artifacts/v2/scripts/run_mlflow_tests.sh

@@ -1,5 +1,7 @@
 #!/bin/bash
 
+set -e
+
 # Run examples for keras, pytorch, sklearn, tensorflow
 cd examples
 

test/test_artifacts/v2/scripts/run_sagemaker_mlflow_tests.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+
+set -e
+
+pytest -s -rA -vv test/unit

test/test_dockerfile_based_harness.py

@@ -85,6 +85,7 @@ def test_dockerfiles_for_cpu(
         ("serve.test.Dockerfile", ["serve-langchain"]),
         ("langchain-aws.test.Dockerfile", ["langchain-aws"]),
         ("mlflow.test.Dockerfile", ["mlflow"]),
+        ("sagemaker-mlflow.test.Dockerfile", ["sagemaker-mlflow"]),
     ],
 )
 def test_dockerfiles_for_gpu(