Ray Integration test fix (pytorch#483)

Summary:
Fixes pytorch#455
* Re-enabled component integration test
* Fixed a regression where Ray Job API changed what their status return function printed
* Added error handling code in `ray_driver.py` that now makes failed scripts, fail - h/t d4l3k
* Fixed broken `CopyComponentProvider` for Ray

Pull Request resolved: pytorch#483

Test Plan:
Ran below locally and everything was green
* `python scripts/component_integration_tests.py --scheduler "ray"`
* `pytest schedulers/test/ray_integration_test.py`

I'll send a followup PR for the other 2 issues https://github.com/pytorch/torchx/issues?q=is%3Aissue+is%3Aopen+label%3Aray

Reviewed By: d4l3k

Differential Revision: D36385069

Pulled By: msaroufim

fbshipit-source-id: aa91b347e99adeea349c0b0fa377a299dfc5613a

Author: msaroufim

.github/workflows/components-integration-tests.yaml

@@ -16,8 +16,7 @@ jobs:
          - scheduler: "kubernetes"
          - scheduler: "local_cwd"
          - scheduler: "local_docker"
-         # TODO uncomment when https://github.com/pytorch/torchx/issues/455 is resolved
-         # - scheduler: "ray"
+         - scheduler: "ray"
       fail-fast: false
     permissions:
       id-token: write

dev-requirements.txt

@@ -23,4 +23,4 @@ ts==0.5.1
 usort==0.6.4
 
 # Ray doesn't support Python 3.10
-ray[default]==1.11.0; python_version < '3.10'
+ray[default]==1.12.0; python_version < '3.10'

scripts/component_integration_tests.py

@@ -100,6 +100,7 @@ def main() -> None:
                 component_provider,
             ],
             "image": torchx_image,
+            "cfg": {},
         },
     }
 

torchx/components/integration_tests/component_provider.py

@@ -104,7 +104,7 @@ def __init__(self, image: str, scheduler: str) -> None:
         self._dst_path = "<None>"
 
     def setUp(self) -> None:
-        if self._scheduler == "local_cwd":
+        if self._scheduler in ["local_cwd", "ray"]:
             fname = "torchx_copy_test.txt"
             self._src_path: str = os.path.join(tempfile.gettempdir(), fname)
             self._dst_path: str = os.path.join(tempfile.gettempdir(), f"{fname}.copy")
@@ -121,7 +121,7 @@ def _process_local_sched(self) -> None:
     def tearDown(self) -> None:
         if os.path.exists(self._dst_path):
             os.remove(self._dst_path)
-        if self._scheduler == "local_cwd" and os.path.exists(self._dst_path):
+        if self._scheduler in ["local_cwd", "ray"] and os.path.exists(self._dst_path):
             os.remove(self._dst_path)
 
     def get_app_def(self) -> AppDef:

torchx/schedulers/ray/ray_driver.py

@@ -49,9 +49,13 @@ def exec_module(self) -> None:
         worker_evn["MASTER_ADDR"] = self.master_addr
         worker_evn["MASTER_PORT"] = str(self.master_port)
         popen = subprocess.Popen(self.cmd, env=worker_evn)
+
         returncode = popen.wait()
         _logger.info(f"Finished with code {returncode}")
 
+        if returncode != 0:
+            raise RuntimeError(f"exec_module failed with return code {returncode}")
+
     def get_actor_address_and_port(self) -> Tuple[str, int]:
         return get_address_and_port()
 
@@ -149,20 +153,14 @@ def main() -> None:  # pragma: no cover
 
     # Await return result of remote ray function
     while len(active_workers) > 0:
+        _logger.info(f"running ray.wait on {active_workers}")
+
         # pyre-fixme[16]: Module `worker` has no attribute `wait`.
         completed_workers, active_workers = ray.wait(active_workers)
         # If a failure occurs the ObjectRef will be marked as completed.
         # Calling ray.get will expose the failure as a RayActorError.
         for object_ref in completed_workers:
-            try:
-                ray.get(object_ref)
-            # If an error occurs during the actor execution,
-            # this error will get propagated as-is to the driver when you call ray.get().
-            # For example, if a ValueError is raised in the actor method call,
-            # this will be raised as a ValueError on the driver.
-            # These exceptions will not be caught in this try-except clause
-            except ray.exceptions.RayActorError as exc:
-                _logger.error("Ray Actor error", exc)
+            ray.get(object_ref)
 
 
 if __name__ == "__main__":

torchx/schedulers/ray_scheduler.py

@@ -308,7 +308,7 @@ def wait_until_finish(self, app_id: str, timeout: int = 30) -> None:
             start = time.time()
             while time.time() - start <= timeout:
                 status_info = client.get_job_status(app_id)
-                status = status_info.status
+                status = status_info
                 if status in {JobStatus.SUCCEEDED, JobStatus.STOPPED, JobStatus.FAILED}:
                     break
                 time.sleep(1)
@@ -322,8 +322,11 @@ def describe(self, app_id: str) -> Optional[DescribeAppResponse]:
             addr, app_id = app_id.split("-")
             client = JobSubmissionClient(f"http://{addr}")
             job_status_info = client.get_job_status(app_id)
-            state = _ray_status_to_torchx_appstate[job_status_info.status]
+            state = _ray_status_to_torchx_appstate[job_status_info]
             roles = [Role(name="ray", num_replicas=-1, image="<N/A>")]
+
+            # get ip_address and put it in hostname
+
             roles_statuses = [
                 RoleStatus(
                     role="ray",
@@ -340,7 +343,7 @@ def describe(self, app_id: str) -> Optional[DescribeAppResponse]:
             return DescribeAppResponse(
                 app_id=app_id,
                 state=state,
-                msg=job_status_info.message or NONE,
+                msg=job_status_info,
                 roles_statuses=roles_statuses,
                 roles=roles,
             )