Skip to content

Commit

Permalink
update log
Browse files Browse the repository at this point in the history
  • Loading branch information
cathyzbn committed Feb 12, 2025
1 parent 241282d commit 713f8fd
Showing 1 changed file with 1 addition and 4 deletions.
5 changes: 1 addition & 4 deletions src/levanter/infra/ray_tpu.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@ def run(self, remote_fn: RemoteFunction | Callable) -> _TpuRunResult:
for i, (actor, info) in enumerate(zip(self.worker_actors, worker_infos)):
if info[0] is None:
raise RuntimeError(f"Worker actor {i} returned invalid info: {info}")
logger.info(f"Initialized worker slice actors {self.worker_actors}")
logger.info(f"Initialized worker slice actors {worker_infos}")
except Exception as e:
self.cancel()
raise RuntimeError("Failed to initialize worker actors") from e
Expand Down Expand Up @@ -208,10 +208,7 @@ def cancel(self) -> None:

def _redecorate_remote_fn_for_tpu(remote_fn, num_hosts, **runtime_env) -> Tuple[RemoteFunction, str]:
if not isinstance(remote_fn, RemoteFunction):
logger.info("CATHY log: decorating non remote function")
remote_fn = ray.remote(remote_fn)
else:
logger.info("CATHY log: decorating remote function")

tpu_name = ray.util.accelerators.tpu.get_current_pod_name() # -> my-tpu
num_tpus_per_host = TPUAcceleratorManager.get_current_node_num_accelerators() # -> 8
Expand Down

0 comments on commit 713f8fd

Please sign in to comment.