Install sagemaker-tensorflow from pypi. Add MKL environment variables for TF 1.8 (#61)

* Switch to install sagemaker-tensorflow from pypi. Add MKL environment variables for TF 1.8.

* Remove sagemaker-tensorflow-extensions from docker file.

Author: nadiaya

docker/1.7.0/final/py2/Dockerfile.cpu

@@ -16,16 +16,10 @@ RUN framework_installable_local=$(basename $framework_installable) && \
     \
     pip install --no-cache --upgrade $framework_installable_local && \
     pip install $framework_support_installable_local && \
+    pip install "sagemaker-tensorflow>=1.7,<1.8" &&\
     \
     rm $framework_installable_local && \
     rm $framework_support_installable_local
 
-COPY sagemaker-tensorflow-extensions sagemaker-tensorflow-extensions
-
-RUN cd sagemaker-tensorflow-extensions && \
-    LD_LIBRARY_PATH=$(dirname $(find /usr/local/lib | grep "libmkl")):$LD_LIBRARY_PATH pip install . && \
-    cd .. && \
-    rm -rf sagemaker-tensorflow-extensions
-
 # entry.py comes from sagemaker-container-support
 ENTRYPOINT ["entry.py"]

docker/1.7.0/final/py2/Dockerfile.gpu

@@ -16,19 +16,11 @@ RUN framework_installable_local=$(basename $framework_installable) && \
     \
     pip install --no-cache --upgrade $framework_installable_local && \
     pip install $framework_support_installable_local && \
+    pip install "sagemaker-tensorflow>=1.7,<1.8" &&\
     \
     rm $framework_installable_local && \
     rm $framework_support_installable_local
 
-COPY sagemaker-tensorflow-extensions sagemaker-tensorflow-extensions
-
-RUN cd sagemaker-tensorflow-extensions && \
-    ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/cuda/lib64/stubs/libcuda.so.1 && \
-    LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH \
-    pip install . && \
-    cd .. && \
-    rm -rf sagemaker-tensorflow-extensions
-
 # entry.py comes from sagemaker-container-support
 ENTRYPOINT ["entry.py"]
 

docker/1.8.0/final/py2/Dockerfile.cpu

@@ -7,20 +7,22 @@ ARG framework_support_installable=sagemaker_tensorflow_container-1.0.0.tar.gz
 
 WORKDIR /root
 
+COPY $framework_installable .
 COPY $framework_support_installable .
 
-RUN framework_support_installable_local=$(basename $framework_support_installable) && \
+RUN framework_installable_local=$(basename $framework_installable) && \
+    framework_support_installable_local=$(basename $framework_support_installable) && \
+    \
+    pip install --no-cache --upgrade $framework_installable_local && \
     pip install $framework_support_installable_local && \
+    pip install "sagemaker-tensorflow>=1.8,<1.9" &&\
+    \
+    rm $framework_installable_local && \
     rm $framework_support_installable_local
 
-RUN pip install --no-cache --upgrade tensorflow==1.8.0
-
-COPY sagemaker-tensorflow-extensions sagemaker-tensorflow-extensions
-
-RUN cd sagemaker-tensorflow-extensions && \
-    pip install . && \
-    cd .. && \
-    rm -rf sagemaker-tensorflow-extensions
+# Set environment variables for MKL
+# TODO: investigate the right value for OMP_NUM_THREADS
+ENV KMP_AFFINITY=granularity=fine,compact,1,0 KMP_BLOCKTIME=1 KMP_SETTINGS=0
 
 # entry.py comes from sagemaker-container-support
 ENTRYPOINT ["entry.py"]

docker/1.8.0/final/py2/Dockerfile.gpu

@@ -16,18 +16,10 @@ RUN framework_installable_local=$(basename $framework_installable) && \
     \
     pip install --no-cache --upgrade $framework_installable_local && \
     pip install $framework_support_installable_local && \
+    pip install "sagemaker-tensorflow>=1.8,<1.9" &&\
     \
     rm $framework_installable_local && \
     rm $framework_support_installable_local
 
-COPY sagemaker-tensorflow-extensions sagemaker-tensorflow-extensions
-
-RUN cd sagemaker-tensorflow-extensions && \
-    ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/cuda/lib64/stubs/libcuda.so.1 && \
-    LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH \
-    pip install . && \
-    cd .. && \
-    rm -rf sagemaker-tensorflow-extensions
-
 # entry.py comes from sagemaker-container-support
 ENTRYPOINT ["entry.py"]