remove PyTorchDistributedDecorator & 2.2.5+zg1.1 (#88)

* remove PyTorchDistributedDecorator * 2.2.5+zg1 * 2.2.5+zg1.1 Co-authored-by: Taleb Zeghmi <[email protected]>
zillow · Jun 10, 2021 · 6101ef0 · 6101ef0
1 parent 7b89959
commit 6101ef0
Show file tree

Hide file tree

Showing 6 changed files with 8 additions and 86 deletions.
diff --git a/metaflow/plugins/__init__.py b/metaflow/plugins/__init__.py
@@ -58,8 +58,7 @@ def _merge_lists(base, overrides, attr):
 from .aws.step_functions.step_functions_decorator import StepFunctionsInternalDecorator
 from .conda.conda_step_decorator import CondaStepDecorator
 from .kfp.kfp_decorator import KfpInternalDecorator
-from .kfp.pytorch_distributed_decorator import PyTorchDistributedDecorator
-from .kfp.accelerator_decorator import AcceleratorDecorator 
+from .kfp.accelerator_decorator import AcceleratorDecorator
 
 STEP_DECORATORS = _merge_lists([CatchDecorator,
                                 TimeoutDecorator,
@@ -69,7 +68,6 @@ def _merge_lists(base, overrides, attr):
                                 BatchDecorator,
                                 StepFunctionsInternalDecorator,
                                 CondaStepDecorator,
-                                PyTorchDistributedDecorator,
                                 AcceleratorDecorator,
                                 KfpInternalDecorator], ext_plugins.STEP_DECORATORS, 'name')
 

diff --git a/metaflow/plugins/kfp/kfp.py b/metaflow/plugins/kfp/kfp.py
@@ -49,7 +49,6 @@
 from .kfp_exit_handler import exit_handler
 from .kfp_foreach_splits import graph_to_task_ids
 from .kfp_get_workflow_uid import get_workflow_uid
-from .pytorch_distributed_decorator import PyTorchDistributedDecorator
 from .accelerator_decorator import AcceleratorDecorator
 from ..aws.batch.batch_decorator import BatchDecorator
 from ..aws.step_functions.schedule_decorator import ScheduleDecorator
@@ -74,7 +73,6 @@ def __init__(
         total_retries: int,
         resource_requirements: Dict[str, str],
         kfp_decorator: KfpInternalDecorator,
-        pytorch_distributed_decorator: PyTorchDistributedDecorator,
         accelerator_decorator: AcceleratorDecorator,
         environment_decorator: EnvironmentDecorator,
     ):
@@ -88,7 +86,6 @@ def __init__(
             if kfp_decorator
             else None
         )
-        self.pytorch_distributed_decorator = pytorch_distributed_decorator
         self.accelerator_decorator = accelerator_decorator
         self.environment_decorator = environment_decorator
 
@@ -391,14 +388,6 @@ def build_kfp_component(node: DAGNode, task_id: str) -> KfpComponent:
                     ),
                     None,  # default
                 ),
-                pytorch_distributed_decorator=next(
-                    (
-                        deco
-                        for deco in node.decorators
-                        if isinstance(deco, PyTorchDistributedDecorator)
-                    ),
-                    None,  # default
-                ),
                 accelerator_decorator=next(
                     (
                         deco
@@ -586,17 +575,9 @@ def _set_container_resources(
             container_op.container.set_ephemeral_storage_limit(
                 resource_requirements["local_storage_limit"]
             )
-        if (
-            kfp_component.pytorch_distributed_decorator
-            or "volume" in resource_requirements
-        ):
-            if kfp_component.pytorch_distributed_decorator:
-                print("This is now deprecated!")
-                mode = [VOLUME_MODE_RWM]
-                volume_dir = "/opt/pytorch_shared/"
-            else:
-                mode = resource_requirements["volume_mode"]
-                volume_dir = resource_requirements["volume_dir"]
+        if "volume" in resource_requirements:
+            mode = resource_requirements["volume_mode"]
+            volume_dir = resource_requirements["volume_dir"]
 
             volume = KubeflowPipelines._create_volume(
                 step_name=kfp_component.name,

diff --git a/metaflow/plugins/kfp/pytorch_distributed_decorator.py b/metaflow/plugins/kfp/pytorch_distributed_decorator.py
diff --git a/metaflow/tutorials/10-pytorch/hello_pytorch.py b/metaflow/tutorials/10-pytorch/hello_pytorch.py
@@ -1,4 +1,4 @@
-from metaflow import FlowSpec, Parameter, step, pytorch_distributed, resources
+from metaflow import FlowSpec, Parameter, step, resources
 
 from models.train import train_model
 from models.evaluate import evaluate_model
@@ -36,8 +36,7 @@ def start(self):
         print(f"ranks: {self.ranks}")
         self.next(self.train, foreach="ranks")
 
-    @resources(cpu=1, cpu_limit=2, gpu="1", memory="2G", memory_limit="5G")
-    @pytorch_distributed
+    @resources(cpu=1, cpu_limit=2, gpu="1", memory="2G", memory_limit="5G", volume="10G")
     @step
     def train(self):
         """

diff --git a/metaflow/tutorials/10-pytorch/models/train.py b/metaflow/tutorials/10-pytorch/models/train.py
@@ -93,7 +93,7 @@ def train_model(
         print("Using distributed PyTorch with {} backend".format(pytorch_backend))
         dist.init_process_group(
             backend=pytorch_backend,
-            init_method="file:///opt/pytorch_shared/sharedfile",
+            init_method="file:///opt/metaflow_volume/sharedfile",
             world_size=world_size,
             rank=rank,
         )

diff --git a/setup.py b/setup.py
@@ -2,7 +2,7 @@
 
 from setuptools import setup, find_packages
 
-version = '2.2.5'
+version = '2.2.5+zg1.1'
 
 # TODO: once this branch is merged or in pip use, remove this
 os.system(