Add API Annotations (#88)

* add annotations

* update test

Author: amogkam

ray_lightning/ray_ddp.py

@@ -13,6 +13,7 @@
 from pytorch_lightning.utilities import rank_zero_only
 
 import ray
+from ray.util import PublicAPI
 from ray.util.queue import Queue
 
 from ray_lightning.session import init_session
@@ -54,6 +55,7 @@ def execute(self, fn: Callable, *args, **kwargs):
         return fn(*args, **kwargs)
 
 
+@PublicAPI(stability="beta")
 class RayPlugin(DDPSpawnPlugin):
     """Pytorch Lightning plugin for DDP training on a Ray cluster.
 

ray_lightning/ray_ddp_sharded.py

@@ -1,7 +1,10 @@
 from pytorch_lightning.plugins import DDPSpawnShardedPlugin
 
+from ray.util import PublicAPI
+
 from ray_lightning import RayPlugin
 
 
+@PublicAPI(stability="beta")
 class RayShardedPlugin(RayPlugin, DDPSpawnShardedPlugin):
     pass

ray_lightning/ray_horovod.py

@@ -5,6 +5,7 @@
 
 import ray
 from ray import ObjectRef
+from ray.util import PublicAPI
 from ray.util.queue import Queue
 
 from ray_lightning.session import init_session
@@ -28,6 +29,7 @@ def get_executable_cls():
     return None
 
 
+@PublicAPI(stability="beta")
 class HorovodRayPlugin(HorovodPlugin):
     """Pytorch Lightning Plugin for Horovod training on a Ray cluster.
 

ray_lightning/tests/test_ddp.py

@@ -1,5 +1,6 @@
 import pytest
 from ray.util.client.ray_client_helpers import ray_start_client_server
+import ray._private.gcs_utils as gcs_utils
 from torch.utils.data import DistributedSampler
 
 from pl_bolts.datamodules import MNISTDataModule
@@ -45,7 +46,7 @@ def check_num_actor():
     plugin = RayPlugin(num_workers=num_workers)
     trainer = get_trainer(tmpdir, plugins=[plugin])
     trainer.fit(model)
-    assert all(actor["State"] == ray.gcs_utils.ActorTableData.DEAD
+    assert all(actor["State"] == gcs_utils.ActorTableData.DEAD
                for actor in list(ray.state.actors().values()))
 
 

ray_lightning/tune.py

@@ -4,6 +4,7 @@
 import os
 
 from pytorch_lightning import Trainer, LightningModule
+from ray.util import PublicAPI
 
 from ray_lightning.session import put_queue, get_actor_rank
 from ray_lightning.util import to_state_stream, Unavailable
@@ -26,6 +27,7 @@ def is_session_enabled():
 
 if TUNE_INSTALLED:
 
+    @PublicAPI(stability="beta")
     def get_tune_ddp_resources(num_workers: int = 1,
                                cpus_per_worker: int = 1,
                                use_gpu: bool = False) -> Dict[str, int]:
@@ -40,6 +42,7 @@ def get_tune_ddp_resources(num_workers: int = 1,
             bundles, strategy="PACK")
         return placement_group_factory
 
+    @PublicAPI(stability="beta")
     class TuneReportCallback(TuneCallback):
         """Distributed PyTorch Lightning to Ray Tune reporting callback
 
@@ -161,6 +164,7 @@ def _handle(self, trainer: Trainer, pl_module: LightningModule):
                 put_queue(lambda: self._create_checkpoint(
                     checkpoint_stream, global_step, self._filename))
 
+    @PublicAPI(stability="beta")
     class TuneReportCheckpointCallback(TuneCallback):
         """PyTorch Lightning to Tune reporting and checkpointing callback.