Lighten logger impact on installation and dynamic import (#285)

* Lighten logger impact in installation and dynamic import

* Simplify imports in itwinai torch trainer and format some files

* Fix undefined name 'ray'

* Fix division by zero

* Return tensor losses

* FIX epoch step return values

env-files/tensorflow/Dockerfile

@@ -41,7 +41,7 @@ ENV PYTHONPATH=""
 # Install itwinai
 COPY pyproject.toml pyproject.toml
 COPY src src
-RUN pip install --no-cache-dir .[nvidia] \
+RUN pip install --no-cache-dir .[prov4ml-nvidia] \
     && itwinai sanity-check --tensorflow --optional-deps ray
 
 # Additional pip deps

env-files/tensorflow/generic_tf.sh

@@ -25,4 +25,4 @@ else
 fi
 
 source $ENV_NAME/bin/activate
-pip install --no-cache-dir -e ".[dev,nvidia,tf]"
+pip install --no-cache-dir -e ".[dev,prov4ml-nvidia,tf]"

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,tf,nvidia,dev] --extra-index-url https://download.pytorch.org/whl/cu124
+RUN pip install --no-cache-dir .[torch,tf,prov4ml-nvidia,dev] --extra-index-url https://download.pytorch.org/whl/cu124
 
 # Install DeepSpeed, Horovod and Ray
 ENV HOROVOD_WITH_PYTORCH=1 \

env-files/torch/generic_torch.sh

@@ -26,6 +26,6 @@ fi
 
 # Activate the venv and then install itwinai as editable
 source $ENV_NAME/bin/activate
-pip install -e ".[torch,tf,dev,nvidia]" \
+pip install -e ".[torch,tf,dev,prov4ml-nvidia]" \
     --no-cache-dir \
     --extra-index-url https://download.pytorch.org/whl/cu121

env-files/torch/jupyter/Dockerfile

@@ -216,7 +216,7 @@ RUN pip install --no-cache-dir --upgrade pip \
 WORKDIR "$HOME/itwinai"
 COPY --chown=${NB_UID} pyproject.toml pyproject.toml
 COPY --chown=${NB_UID} 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,prov4ml-nvidia,dev] --extra-index-url https://download.pytorch.org/whl/cu124
 
 # Apex: https://github.com/NVIDIA/apex
 # (needed for DeepSpeed *_FUSED optinal build options)

env-files/torch/slim.Dockerfile

@@ -56,7 +56,7 @@ RUN /usr/bin/python3.10 -m venv /opt/venv \
 WORKDIR /app
 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,prov4ml-nvidia,dev] --extra-index-url https://download.pytorch.org/whl/cu124
 
 # Install DeepSpeed, Horovod and Ray
 RUN CONTAINER_TORCH_VERSION="$(python -c 'import torch;print(torch.__version__)')" \

pyproject.toml

@@ -40,9 +40,8 @@ dependencies = [
   "jsonargparse[signatures]>=4.34.0",
   "matplotlib>=3.9.2",
   "pip>=24.3.1",
-  "prov4ml@git+https://github.com/matbun/ProvML@new-main",
   "ray[default,train,tune]>=2.39.0",
-  "prov>=2.0.1",
+  "tensorboard>=2.16.2",
 ]
 
 [project.optional-dependencies]
@@ -79,10 +78,12 @@ hpo = [
   "hpbandster>=0.7.0",
   "gpy>=1.13.2",
 ]
-
-macos = ["prov4ml[apple]@git+https://github.com/matbun/ProvML@new-main"]
-nvidia = ["prov4ml[nvidia]@git+https://github.com/matbun/ProvML@new-main"]
-# amd = ["prov4ml[amd]@git+https://github.com/matbun/ProvML@new-main"]
+prov4ml = ["prov4ml@git+https://github.com/matbun/ProvML@new-main"]
+prov4ml-macos = ["prov4ml[apple]@git+https://github.com/matbun/ProvML@new-main"]
+prov4ml-nvidia = [
+  "prov4ml[nvidia]@git+https://github.com/matbun/ProvML@new-main",
+]
+# prov4ml-amd = ["prov4ml[amd]@git+https://github.com/matbun/ProvML@new-main"]
 
 [project.urls]
 Homepage = "https://www.intertwin.eu/"