Remove resource limit from pytorch job (#466)

* Remove resource limit from pytorch job Signed-off-by: Yuvraj <code@evalsocket.dev>
flyteorg · Oct 29, 2021 · a653b22 · a653b22
1 parent 69ef957
commit a653b22
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 6 deletions.
diff --git a/cookbook/integrations/kubernetes/kfmpi/requirements.in b/cookbook/integrations/kubernetes/kfmpi/requirements.in
@@ -1,4 +1,4 @@
 flytekitplugins-kfmpi>=0.23.2
-flytekit>=0.23.2
+flytekit>=0.24.0b0
 flyteidl>=0.21.4
 tensorflow==2.6.0
diff --git a/cookbook/integrations/kubernetes/kfmpi/requirements.txt b/cookbook/integrations/kubernetes/kfmpi/requirements.txt
@@ -22,6 +22,8 @@ clang==5.0
     # via tensorflow
 click==7.1.2
     # via flytekit
+cloudpickle==2.0.0
+    # via flytekit
 croniter==1.0.15
     # via flytekit
 dataclasses-json==0.5.6
@@ -38,12 +40,12 @@ docstring-parser==0.12
     # via flytekit
 flatbuffers==1.12
     # via tensorflow
-flyteidl==0.21.7
+flyteidl==0.21.8
     # via
     #   -r requirements.in
     #   flytekit
     #   flytekitplugins-kfmpi
-flytekit==0.23.2
+flytekit==0.24.0b0
     # via
     #   -r requirements.in
     #   flytekitplugins-kfmpi
@@ -106,7 +108,7 @@ opt-einsum==3.3.0
     # via tensorflow
 pandas==1.3.4
     # via flytekit
-protobuf==3.19.0
+protobuf==3.19.1
     # via
     #   flyteidl
     #   flytekit

diff --git a/cookbook/integrations/kubernetes/kfpytorch/pytorch_mnist.py b/cookbook/integrations/kubernetes/kfpytorch/pytorch_mnist.py
@@ -180,12 +180,12 @@ class Hyperparameters(object):
 @task(
     task_config=PyTorch(
         num_workers=2,
-        per_replica_requests=Resources(cpu=cpu_request, mem=mem_request, gpu=gpu_request),
-        per_replica_limits=Resources(mem=mem_limit, gpu=gpu_limit),
     ),
     retries=2,
     cache=True,
     cache_version="1.0",
+    requests=Resources(cpu=cpu_request, mem=mem_request, gpu=gpu_request),
+    limits=Resources(mem=mem_limit, gpu=gpu_limit)
 )
 def mnist_pytorch_job(hp: Hyperparameters) -> TrainingOutputs:
     log_dir = "logs"