improve class and variable naming

Signed-off-by: Andrew Sy Kim <andrewsy@google.com>

Author: andrewsykim

skyrl/backends/jax.py

@@ -113,6 +113,7 @@ class JaxBackendConfig(BaseModel, extra="forbid"):
         default=None,
         description="Total number of processes in the multi-node cluster",
     )
+    # RayJaxBackend configuration
     ray_actor_options: dict = Field(
         default_factory=dict,
         description="Options to pass to Ray actors (e.g., resources, num_cpus)",

skyrl/backends/ray_jax.py

@@ -1,4 +1,3 @@
-import socket
 import ray
 from cloudpathlib import AnyPath
 
@@ -7,18 +6,12 @@
 from skyrl.tinker import types
 from skyrl.utils.log import logger
 
-def get_free_port() -> int:
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        s.bind(("", 0))
-        return s.getsockname()[1]
-
 
 @ray.remote
-class RayJaxActor:
-    """Ray Actor wrapper for JaxBackendImpl.
-    
-    Each actor runs JaxBackendImpl and communicates with other actors
-    via JAX distributed (NCCL) for data parallel operations.
+class RayJaxBackendImpl:
+    """RayJaxBackendImpl is a Ray wrapper for JaxBackendImpl.
+
+    Each actor calls jax.distributed.initialize() and holds an instance of JaxBackendImpl.
     """
     def __init__(self, base_model: str, config: JaxBackendConfig, process_id: int):
         self.base_model = base_model
@@ -28,7 +21,7 @@ def __init__(self, base_model: str, config: JaxBackendConfig, process_id: int):
 
         if process_id == 0:
             self.node_ip = ray.util.get_node_ip_address()
-            self.port = get_free_port()
+            self.port = 7777
             self.coordinator_address = f"{self.node_ip}:{self.port}"
         else:
             self.coordinator_address = None
@@ -38,7 +31,7 @@ def get_coordinator_address(self) -> str:
 
     def setup(self, coordinator_address: str | None = None):
         """Initializes JAX distributed and creates JaxBackendImpl."""
-        import jax # Import here to avoid issues or ensure it's loaded in the actor
+        import jax
 
         addr = coordinator_address or self.coordinator_address
         logger.info(f"Worker {self.process_id} initializing JAX distributed with coordinator {addr}")
@@ -51,10 +44,6 @@ def setup(self, coordinator_address: str | None = None):
         self.backend = JaxBackendImpl(self.base_model, self.config, self.process_id)
         logger.info(f"Worker {self.process_id} JaxBackendImpl initialized.")
 
-    # =========================================================================
-    # Proxied Backend Methods
-    # =========================================================================
-
     def create_model(self, model_id: str, lora_config: types.LoraConfig) -> None:
         self.backend.create_model(model_id, lora_config)
 
@@ -90,61 +79,61 @@ def get_metrics(self) -> types.EngineMetrics:
 
 
 class RayJaxBackend(AbstractBackend):
-    """Proxy Backend that orchestrates Ray actors for multi-node JAX execution.
-    
-    This class runs in the driver program (Tinker Engine process) and proxies
+    """RayJaxBackend is a proxy Backend that orchestrates Ray actors for multi-node JAX execution.
+
+    This class runs in the driver program (along with Tinker API / Engine) and proxies
     commands to all JAX workers running as Ray actors.
     """
     def __init__(self, base_model: str, config: JaxBackendConfig):
         self.base_model = base_model
         self.config = config.model_copy()
-        
+
         num_processes = self.config.num_processes or 1
         self.config.num_processes = num_processes
-        
+
         logger.info(f"Initializing RayJaxBackend with num_processes={num_processes}")
 
-        # Instantiate a Ray placement group
+        # Initialize a Ray placement group based on ray_pg_bundles in JaxBackendConfig
         from ray.util.placement_group import placement_group
-        logger.info("Instantiating Ray placement group for JAX workers...")
+        logger.info("Creating Ray placement group for JAX backend")
         bundles = self.config.ray_pg_bundles
         if not bundles:
             bundles = [{"CPU": 1}] * num_processes
         self.pg = placement_group(bundles, strategy="SPREAD")
         ray.get(self.pg.ready())
 
         self.workers = []
-        
-        # Create worker 0 (coordinator)
+
+        # node0 (coordinator)
+        logger.info("Scheduling Ray actor for node0 (JAX coordinator)")
         w0_options = self.config.ray_actor_options.copy()
         w0_options.update({
             "placement_group": self.pg,
             "placement_group_bundle_index": 0,
         })
-        w0 = RayJaxActor.options(**w0_options).remote(self.base_model, self.config, 0)
+        w0 = RayJaxBackendImpl.options(**w0_options).remote(self.base_model, self.config, 0)
         self.workers.append(w0)
 
-        # Retrieve dynamically allocated coordinator address from actor 0
         coordinator_address = ray.get(w0.get_coordinator_address.remote())
-        
-        # Create other workers
+
+        # Create remaining node1 - nodeN for multi-node training.
+        logger.info("Scheduling remaining Ray actors (JAX workers)")
         for i in range(1, num_processes):
             wi_options = self.config.ray_actor_options.copy()
             wi_options.update({
                 "placement_group": self.pg,
                 "placement_group_bundle_index": i,
             })
-            w = RayJaxActor.options(**wi_options).remote(self.base_model, self.config, i)
+            w = RayJaxBackendImpl.options(**wi_options).remote(self.base_model, self.config, i)
             self.workers.append(w)
 
-        # Trigger setup on all workers
         # This will block until JAX distributed is initialized on all workers
         setup_refs = [w0.setup.remote()]
         for w in self.workers[1:]:
             setup_refs.append(w.setup.remote(coordinator_address))
 
         ray.get(setup_refs)
-        logger.info("RayJaxBackend is fully initialized and distributed cluster is ready.")
+        logger.info("RayJaxBackend is fully initialized and distributed JAX cluster is ready.")
 
     @property
     def metrics(self) -> types.EngineMetrics: