Skip to content

Commit 936016c

Browse files
authored
[checkpointing] support ipv6 (#12)
1 parent aec61c7 commit 936016c

File tree

1 file changed

+6
-1
lines changed

1 file changed

+6
-1
lines changed

torchft/checkpointing.py

+6-1
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,11 @@
1616
logger: logging.Logger = logging.getLogger(__name__)
1717

1818

19+
class IPv6HTTPServer(ThreadingHTTPServer):
20+
address_family = socket.AF_INET6
21+
request_queue_size = 1024
22+
23+
1924
class CheckpointServer:
2025
def __init__(self, state_dict) -> None:
2126
self._checkpoint_lock = threading.Lock()
@@ -53,7 +58,7 @@ def err(self, msg: str) -> None:
5358
self.wfile.write(msg.encode())
5459

5560
server_address = ("", 0)
56-
self._server = ThreadingHTTPServer(server_address, RequestHandler)
61+
self._server = IPv6HTTPServer(server_address, RequestHandler)
5762
logger.info(f"Started CheckpointServer on {self.address()}...")
5863

5964
self._thread = threading.Thread(

0 commit comments

Comments
 (0)