Revert "Moved cpu processing into pre and postprocessing (kubeflow#83)…

…" (kubeflow#85) This reverts commit cadcc11.
georgkaleido · Dec 20, 2021 · c0f3582 · c0f3582
1 parent dd5630e
commit c0f3582
Show file tree

Hide file tree

Showing 6 changed files with 46 additions and 50 deletions.
diff --git a/core/bin/removebg-server.py b/core/bin/removebg-server.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python
 import math
+import multiprocessing
 import os
 import time
 from dataclasses import dataclass
@@ -124,8 +125,6 @@ def _preprocess(self, processing_data: ImageProcessingData, message: Dict[bytes,
             f"{image.width}x{image.height} -> {im_cv.shape[1]}x{im_cv.shape[0]} (crop: {crop})",
         )
 
-        im_cv = to_tensor(im_cv, bgr2rgb=True, cuda=False)
-
         # pass all data which is needed for postprocessing
         processing_data.data = {
             "times": times,
@@ -150,16 +149,8 @@ def _post_process(self, processing_data: ImageProcessingData) -> None:
         result = processing_data.result
         data = processing_data.data
 
-        # when it is not mocked
-        if data["data"] is None:
-            im_tr_rgb = data["im_tr_rgb"]
-            im_tr_alpha = data["im_tr_alpha"]
-            im_rgb = to_numpy(im_tr_rgb, rgb2bgr=True, cpu_clamp=True)
-            im_alpha = to_numpy(im_tr_alpha)
-            im_res = np.dstack((im_rgb, np.expand_dims(im_alpha, axis=2)))
-        else:
-            im_res = data["data"]
-
+        # load data from preprocessing
+        im_res = data["data"]
         times = data["times"]
         image = data["image"]
         semitransparency = data["semitransparency"]
@@ -325,6 +316,7 @@ def __init__(
         mock_response: bool,
         require_models: bool,
         worker_init_kwargs: Optional[Dict[str, Any]] = None,
+        worker_count: int = multiprocessing.cpu_count(),
     ) -> None:
         super().__init__(
             request_queue,
@@ -334,6 +326,7 @@ def __init__(
             rabbitmq_password,
             worker_class=worker_class,
             worker_init_kwargs=worker_init_kwargs,
+            worker_count=worker_count,
         )
         self.removebg = None
         self.identifier = None
@@ -367,7 +360,7 @@ def _process(self, data: ImageProcessingData) -> None:
             processing_data["api"] = "mock"
         else:
             # transfer to gpu
-            im_tr = im.cuda(non_blocking=True)
+            im_tr = to_tensor(im, bgr2rgb=True)
 
             # overwrite if auto
             if processing_data["api"] == "auto":
@@ -389,9 +382,10 @@ def _process(self, data: ImageProcessingData) -> None:
                 result.status = ERROR_STATUS
                 result.description = UNKNOWN_FOREGROUND
             else:
-                # move model results to cpu
-                processing_data["im_tr_rgb"] = im_tr_rgb.cpu()
-                processing_data["im_tr_alpha"] = im_tr_alpha.cpu()
+                im_rgb = to_numpy(im_tr_rgb, rgb2bgr=True)
+                im_alpha = to_numpy(im_tr_alpha)
+
+                processing_data["data"] = np.dstack((im_rgb, np.expand_dims(im_alpha, axis=2)))
 
         processing_data["time"] = time.time() - t
 
@@ -400,12 +394,19 @@ def main():
     rabbitmq_args = read_rabbitmq_env_variables()
     mock_response = bool(int(os.environ.get("MOCK_RESPONSE", 0)))
     require_models = bool(int(os.environ.get("REQUIRE_MODELS", 1)))
+    # number of workers to spawn, defaults to number of cpus
+    # be aware that the default worker count might include
+    # "virtual" hyperthreaded cpus and impacts memory usage
+    worker_count = min(
+        int(os.environ.get("MAX_WORKER_COUNT", multiprocessing.cpu_count())), multiprocessing.cpu_count()
+    )
 
     server = RemovebgServer(
         *rabbitmq_args,
         require_models=require_models,
         worker_class=RemovebgWorker,
         mock_response=mock_response,
+        worker_count=worker_count,
     )
     server.start()
 

diff --git a/core/requirements-deploy.txt b/core/requirements-deploy.txt
@@ -8,11 +8,11 @@ certifi==2021.10.8
     # via
     #   requests
     #   sentry-sdk
-charset-normalizer==2.0.9
+charset-normalizer==2.0.8
     # via requests
 click==8.0.3
     # via wandb
-configparser==5.2.0
+configparser==5.1.0
     # via wandb
 docker-pycreds==0.4.0
     # via wandb
@@ -24,7 +24,7 @@ google==3.0.0
     # via kaleido-core
 idna==3.3
     # via requests
-kaleido-core==2.9.0
+kaleido-core==2.8.0
     # via -r requirements-deploy.in
 msgpack==0.6.2
     # via
@@ -45,7 +45,7 @@ opencv-python==4.5.4.60
     # via
     #   -r requirements-deploy.in
     #   kaleido-core
-pandas==1.3.5
+pandas==1.3.4
     # via -r requirements-deploy.in
 pathtools==0.1.2
     # via wandb

diff --git a/core/requirements-dev.txt b/core/requirements-dev.txt
@@ -24,19 +24,19 @@ cffi==1.15.0
     # via cryptography
 cfgv==3.3.1
     # via pre-commit
-charset-normalizer==2.0.9
+charset-normalizer==2.0.8
     # via requests
 click==8.0.3
     # via
     #   black
     #   wandb
 colorama==0.4.4
     # via twine
-configparser==5.2.0
+configparser==5.1.0
     # via wandb
 cryptography==36.0.0
     # via secretstorage
-distlib==0.3.4
+distlib==0.3.3
     # via virtualenv
 docker-pycreds==0.4.0
     # via wandb
@@ -68,7 +68,7 @@ jeepney==0.7.1
     # via
     #   keyring
     #   secretstorage
-kaleido-core==2.9.0
+kaleido-core==2.8.0
     # via -r requirements-deploy.in
 keyring==23.4.0
     # via twine
@@ -102,7 +102,7 @@ packaging==21.3
     #   bleach
     #   build
     #   pytest
-pandas==1.3.5
+pandas==1.3.4
     # via -r requirements-deploy.in
 pathspec==0.9.0
     # via black
@@ -157,7 +157,7 @@ pyyaml==6.0
     # via
     #   pre-commit
     #   wandb
-readme-renderer==31.0
+readme-renderer==30.0
     # via twine
 regex==2021.11.10
     # via black
@@ -218,7 +218,7 @@ tqdm==4.62.3
     # via
     #   kaleido-core
     #   twine
-twine==3.7.1
+twine==3.7.0
     # via -r requirements-dev.in
 typing-extensions==4.0.1
     # via

diff --git a/core/requirements-eval.txt b/core/requirements-eval.txt
@@ -18,13 +18,13 @@ certifi==2021.10.8
     # via
     #   requests
     #   sentry-sdk
-charset-normalizer==2.0.9
+charset-normalizer==2.0.8
     # via
     #   aiohttp
     #   requests
 click==8.0.3
     # via wandb
-configparser==5.2.0
+configparser==5.1.0
     # via wandb
 crcmod==1.7
     # via -r requirements-eval.in
@@ -44,12 +44,12 @@ gitpython==3.1.24
     # via wandb
 google==3.0.0
     # via kaleido-core
-google-api-core==2.3.0
+google-api-core==2.2.2
     # via
     #   google-api-python-client
     #   google-cloud-core
     #   google-cloud-storage
-google-api-python-client==2.33.0
+google-api-python-client==2.32.0
     # via -r requirements-eval.in
 google-auth==2.3.3
     # via
@@ -78,7 +78,7 @@ google-crc32c==1.3.0
     #   google-resumable-media
 google-resumable-media==2.1.0
     # via google-cloud-storage
-googleapis-common-protos==1.54.0
+googleapis-common-protos==1.53.0
     # via google-api-core
 grpcio==1.42.0
     # via tensorboard
@@ -92,7 +92,7 @@ idna==3.3
     #   yarl
 importlib-metadata==4.8.2
     # via markdown
-kaleido-core==2.9.0
+kaleido-core==2.8.0
     # via -r requirements.in
 markdown==3.3.6
     # via tensorboard
@@ -134,7 +134,7 @@ packaging==21.3
     # via
     #   pytorch-lightning
     #   torchmetrics
-pandas==1.3.5
+pandas==1.3.4
     # via
     #   -r requirements-eval.in
     #   -r requirements.in
@@ -179,7 +179,7 @@ python-dateutil==2.8.2
     # via
     #   pandas
     #   wandb
-pytorch-lightning==1.5.5
+pytorch-lightning==1.5.4
     # via -r requirements-eval.in
 pytz==2021.3
     # via pandas
@@ -240,7 +240,7 @@ tensorboard-plugin-wit==1.8.0
     #   torchmetrics
     #   torchtext
     #   torchvision
-torchmetrics==0.6.1
+torchmetrics==0.6.0
     # via pytorch-lightning
     # via -r requirements-eval.in
     # via

diff --git a/core/requirements-train.txt b/core/requirements-train.txt
@@ -18,13 +18,13 @@ certifi==2021.10.8
     # via
     #   requests
     #   sentry-sdk
-charset-normalizer==2.0.9
+charset-normalizer==2.0.8
     # via
     #   aiohttp
     #   requests
 click==8.0.3
     # via wandb
-configparser==5.2.0
+configparser==5.1.0
     # via wandb
 crcmod==1.7
     # via -r requirements-train.in
@@ -46,12 +46,12 @@ gitpython==3.1.24
     # via wandb
 google==3.0.0
     # via kaleido-core
-google-api-core==2.3.0
+google-api-core==2.2.2
     # via
     #   google-api-python-client
     #   google-cloud-core
     #   google-cloud-storage
-google-api-python-client==2.33.0
+google-api-python-client==2.32.0
     # via -r requirements-train.in
 google-auth==2.3.3
     # via
@@ -80,7 +80,7 @@ google-crc32c==1.3.0
     #   google-resumable-media
 google-resumable-media==2.1.0
     # via google-cloud-storage
-googleapis-common-protos==1.54.0
+googleapis-common-protos==1.53.0
     # via google-api-core
 grpcio==1.42.0
     # via tensorboard
@@ -94,7 +94,7 @@ idna==3.3
     #   yarl
 importlib-metadata==4.8.2
     # via markdown
-kaleido-core==2.9.0
+kaleido-core==2.8.0
     # via -r requirements.in
 markdown==3.3.6
     # via tensorboard
@@ -132,7 +132,7 @@ packaging==21.3
     # via
     #   pytorch-lightning
     #   torchmetrics
-pandas==1.3.5
+pandas==1.3.4
     # via -r requirements.in
 pathtools==0.1.2
     # via wandb
@@ -176,7 +176,7 @@ python-dateutil==2.8.2
     #   wandb
 python-etcd==0.4.5
     # via torchelastic
-pytorch-lightning==1.5.5
+pytorch-lightning==1.5.4
     # via -r requirements-train.in
 pytz==2021.3
     # via pandas
@@ -238,7 +238,7 @@ tensorboard-plugin-wit==1.8.0
     #   torchvision
 torchelastic==0.2.0
     # via -r requirements-train.in
-torchmetrics==0.6.1
+torchmetrics==0.6.0
     # via pytorch-lightning
     # via
     #   -r requirements.in

diff --git a/docker-compose.yml b/docker-compose.yml
@@ -60,11 +60,6 @@ services:
 
   # Core
   core:
-    # Default shared memory segment for pytorch multiprocessing not enough
-    # solved by setting ipc host.
-    # see https://discuss.pytorch.org/t/unable-to-write-to-file-torch-18692-1954506624/9990/2
-    ipc: host
-
     build:
       context: core
       dockerfile: Dockerfile