Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Improve performance by increase recv size on service socket #4122

Merged
merged 5 commits into from Aug 18, 2022
Merged

Improve performance by increase recv size on service socket #4122

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
c66f552
merge
raubitsj Aug 17, 2022
ea87c82
remove sockbuffer
raubitsj Aug 17, 2022
aeb9f64
Merge remote-tracking branch 'origin/master' into perf-increase-recv-…
raubitsj Aug 18, 2022
a689bd2
drop to 64k
raubitsj Aug 18, 2022
2640b77
Merge branch 'master' into perf-increase-recv-size
raubitsj Aug 18, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
11 changes: 10 additions & 1 deletion wandb/sdk/lib/sock_client.py
View file
Open in desktop
Expand Up @@ -83,6 +83,7 @@ class SockClient:
_sockid: str
_retry_delay: float
_lock: "threading.Lock"
_bufsize: int
_buffer: SockBuffer

# current header is magic byte "W" followed by 4 byte length of the message
Expand All @@ -93,12 +94,19 @@ def __init__(self) -> None:
self._sockid = uuid.uuid4().hex
self._retry_delay = 0.1
self._lock = threading.Lock()
self._bufsize = 4096
self._buffer = SockBuffer()

def connect(self, port: int) -> None:
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.connect(("localhost", port))
self._sock = s
self._detect_bufsize()

def _detect_bufsize(self) -> None:
sndbuf_size = self._sock.getsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF)
rcvbuf_size = self._sock.getsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF)
self._bufsize = min(sndbuf_size, rcvbuf_size, 65536)

def close(self) -> None:
self._sock.close()
Expand All @@ -108,6 +116,7 @@ def shutdown(self, val: int) -> None:

def set_socket(self, sock: socket.socket) -> None:
self._sock = sock
self._detect_bufsize()

def _sendall_with_error_handle(self, data: bytes) -> None:
# This is a helper function for sending data in a retry fashion.
Expand Down Expand Up @@ -243,7 +252,7 @@ def _read_packet_bytes(self, timeout: int = None) -> Optional[bytes]:
if timeout:
self._sock.settimeout(timeout)
try:
data = self._sock.recv(4096)
data = self._sock.recv(self._bufsize)
except socket.timeout:
break
except ConnectionResetError:
Expand Down