[autoscaler] Optimize finding the node id (ray-project#21885)

This is a simple refactoring change and my first PR in ray-project. This change moves an if statement outside of a loop. This way the check is not repeated for each iteration.

python/ray/autoscaler/node_provider.py

@@ -44,13 +44,15 @@ def non_terminated_nodes(self, tag_filters: Dict[str, str]) -> List[str]:
         """Return a list of node ids filtered by the specified tags dict.
 
         This list must not include terminated nodes. For performance reasons,
-        providers are allowed to cache the result of a call to nodes() to
-        serve single-node queries (e.g. is_running(node_id)). This means that
-        nodes() must be called again to refresh results.
+        providers are allowed to cache the result of a call to
+        non_terminated_nodes() to serve single-node queries
+        (e.g. is_running(node_id)). This means that non_terminate_nodes() must
+        be called again to refresh results.
 
         Examples:
             >>> provider.non_terminated_nodes({TAG_RAY_NODE_KIND: "worker"})
             ["node-1", "node-2"]
+
         """
         raise NotImplementedError
 
@@ -97,13 +99,11 @@ def find_node_id():
 
         if not find_node_id():
             all_nodes = self.non_terminated_nodes({})
+            ip_func = self.internal_ip if use_internal_ip else self.external_ip
+            ip_cache = self._internal_ip_cache if use_internal_ip \
+                else self._external_ip_cache
             for node_id in all_nodes:
-                if use_internal_ip:
-                    int_ip = self.internal_ip(node_id)
-                    self._internal_ip_cache[int_ip] = node_id
-                else:
-                    ext_ip = self.external_ip(node_id)
-                    self._external_ip_cache[ext_ip] = node_id
+                ip_cache[ip_func(node_id)] = node_id
 
         if not find_node_id():
             if use_internal_ip: