fix docker image

Author: priyakasimbeg

docker/Dockerfile

@@ -23,6 +23,8 @@ RUN apt-get update && apt-get install -y \
     libreadline-dev \
     libffi-dev \
     curl \
+    liblzma-dev \
+    libbz2-dev \
     vim
 
 # Download and install Python 3.11

docker/build_docker_images.sh

@@ -1,34 +1,40 @@
 #!/bin/bash
 # Bash script to build and push dev docker images to artifact repo
 # Usage:
-#   bash build_docker_images.sh -b <git_branch>
+#   bash build_docker_images.sh -b <git_branch> -f <framework>
 
 # Make program exit with non-zero exit code if any command fails.
 set -e
 
-while getopts b: flag
+while getopts "b:p:f:" flag;
 do
     case "${flag}" in
         b) GIT_BRANCH=${OPTARG};;
-    esac
-    case "${flag}" in 
         p) PROJECT=${OPTARG};;
+        f) FRAMEWORK=${OPTARG};;
     esac
 done
 
 # Artifact repostiory
-if [ "$PROJECT" = "mlcommons-algoperf"]; then
+if [ "$PROJECT" = "mlcommons-algoperf" ]; then
     ARTIFACT_REPO="europe-west-4-docker.pkg.dev/mlcommons-algoperf/algoperf-docker-repo"
 else
     ARTIFACT_REPO="us-central1-docker.pkg.dev/training-algorithms-external/mlcommons-docker-repo"
 fi
 
-if [[ -z ${GIT_BRANCH+x} ]]
+if [[ -z ${GIT_BRANCH+x} ]];
 then 
 GIT_BRANCH='main' # Set default argument
 fi 
 
-for FRAMEWORK in "jax" "pytorch" "both"
+FRAMEWORKS=( "jax" "pythorch" "both" )
+
+if [[ -n "$FRAMEWORK" ]];
+then
+    FRAMEWORKS=("$FRAMEWORK")
+fi
+
+for FRAMEWORK in "${FRAMEWORKS[@]}";
 do
     IMAGE_NAME="algoperf_${FRAMEWORK}_${GIT_BRANCH}"
     DOCKER_BUILD_COMMAND="docker build --no-cache -t $IMAGE_NAME . --build-arg framework=$FRAMEWORK --build-arg branch=$GIT_BRANCH"

pyproject.toml

@@ -39,7 +39,7 @@ dependencies = [
   "numpy>=2.0.2",
   "pandas>=2.0.1",
   "tensorflow==2.19.0",
-  "tensorflow-datasets==4.9.7",
+  "tensorflow-datasets==4.9.9",
   "tensorflow-probability==0.20.0",
   "tensorflow-addons==0.20.0",
   "gputil==1.4.0",