Add support for sagemaker-headless-execution-driver (#105)

* Add support for sagemaker-headless-execution-driver * fix import issue * Bug fix: Removed extra comma --------- Co-authored-by: Mustafa Oztoprak <[email protected]>
aws · Oct 24, 2023 · 593e3e7 · 593e3e7
1 parent 45f0b7d
commit 593e3e7
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 0 deletions.
diff --git a/build_artifacts/v0/v0.9/v0.9.0/cpu.env.in b/build_artifacts/v0/v0.9/v0.9.0/cpu.env.in
@@ -9,6 +9,7 @@ conda-forge::pandas[version='>=2.0.3,<3.0.0']
 conda-forge::scikit-learn[version='>=1.3.1,<2.0.0']
 conda-forge::jinja2[version='>=3.1.2,<4.0.0']
 conda-forge::matplotlib[version='>=3.7.3,<4.0.0']
+conda-forge::sagemaker-headless-execution-driver[version'>=0.0.9']
 conda-forge::ipython[version='>=8.12.2,<9.0.0']
 conda-forge::scipy[version='>=1.10.1,<2.0.0']
 conda-forge::keras[version='>=2.12.0,<3.0.0']

diff --git a/build_artifacts/v0/v0.9/v0.9.0/gpu.env.in b/build_artifacts/v0/v0.9/v0.9.0/gpu.env.in
@@ -9,6 +9,7 @@ conda-forge::pandas[version='>=2.0.3,<3.0.0']
 conda-forge::scikit-learn[version='>=1.3.1,<2.0.0']
 conda-forge::jinja2[version='>=3.1.2,<4.0.0']
 conda-forge::matplotlib[version='>=3.7.3,<4.0.0']
+conda-forge::sagemaker-headless-execution-driver[version'>=0.0.9']
 conda-forge::ipython[version='>=8.12.2,<9.0.0']
 conda-forge::scipy[version='>=1.10.1,<2.0.0']
 conda-forge::keras[version='>=2.12.0,<3.0.0']

diff --git a/test/test_artifacts/v0/sagemaker-headless-execution-driver.test.Dockerfile b/test/test_artifacts/v0/sagemaker-headless-execution-driver.test.Dockerfile
@@ -0,0 +1,8 @@
+ARG COSMOS_IMAGE
+FROM $COSMOS_IMAGE
+
+ARG MAMBA_DOCKERFILE_ACTIVATE=1
+
+
+# Execute execution_driver module to for sagemaker-headless-execution-driver installation
+CMD ["python", "-c",  "import sagemaker_headless_execution_driver.headless_execution as execution_driver"]