Update Dockerfile (#273)

ci/dagger.json

@@ -1,6 +1,6 @@
 {
   "name": "itwinai",
+  "engineVersion": "v0.14.0",
   "sdk": "python",
-  "source": ".",
-  "engineVersion": "v0.13.6"
+  "source": "."
 }

env-files/torch/Dockerfile

@@ -35,7 +35,7 @@ RUN pip install --no-cache-dir --upgrade pip \
 # Install itwinai with torch
 COPY pyproject.toml pyproject.toml
 COPY src src
-RUN pip install --no-cache-dir .[torch,nvidia,dev] --extra-index-url https://download.pytorch.org/whl/cu124
+RUN pip install --no-cache-dir .[torch,tf,nvidia,dev] --extra-index-url https://download.pytorch.org/whl/cu124
 
 # Install DeepSpeed, Horovod and Ray
 ENV HOROVOD_WITH_PYTORCH=1 \
@@ -62,16 +62,20 @@ RUN CONTAINER_TORCH_VERSION="$(python -c 'import torch;print(torch.__version__)'
     && pip install --no-cache-dir torch=="$CONTAINER_TORCH_VERSION" \
     deepspeed==0.15.* \
     git+https://github.com/horovod/horovod.git@3a31d93 \
-    ray[tune] \
     # fix .triton/autotune/Fp16Matmul_2d_kernel.pickle bug
     && pver="$(python --version 2>&1 | awk '{print $2}' | cut -f1-2 -d.)" \
     && line=$(cat -n "/usr/local/lib/python${pver}/dist-packages/deepspeed/ops/transformer/inference/triton/matmul_ext.py" | grep os.rename | awk '{print $1}' | head -n 1) \
     && sed -i "${line}s|^|#|" "/usr/local/lib/python${pver}/dist-packages/deepspeed/ops/transformer/inference/triton/matmul_ext.py"
 
 # Installation sanity check
 RUN itwinai sanity-check --torch \
+    # Needed by TensorBoard logger
+    --optional-deps tensorflow \
+    # Microsoft DeepSpeed for distributed ML
     --optional-deps deepspeed \
+    # Horovod for distributed ML
     --optional-deps horovod \
+    # Ray for disrtibuted ML and hyperparameter-tuning
     --optional-deps ray
 
 # Additional pip deps