[WB-6939] Add service backend using sockets (support fork) by raubitsj · Pull Request #2892 · wandb/wandb

raubitsj · 2021-11-07T17:28:11Z

Implements: WB-6939

Description

Changes the default "wandb-service" mode to use raw tcp sockets instead of GRPC.

This is a desirable change because:

gRPC is not fork safe (forking a client grpc process after a connection has been made is unsafe)
gRPC introduces overhead and complexity that is unnecessary in many same node cases

This change (introducing a new transport method) is undesirable because:

It requires maintaining a separate code path
Getting socket error handling and shutdown handling can be difficult to get right

Additional changes:

Converts type annotation to be opt-out instead of opt-in (for wandb/sdk/ only right now)

TODO:

Remove debug messages
Fix up socket shutdown
Move grpc server code to its own place
~~Disable RunStatusChecker for now~~ (works now)
Add grpc specific tests
Make tests pass!
Add (update) dev docs -- sequence diagrams
Make sure exception handling works
Improve coverage

Future:

Enable RunStatusChecker again (WB-7352)
windows tests / support

Testing

functional tests only (more unittests can come later once things stabilize)

codecov · 2021-11-07T17:41:59Z

Codecov Report

Merging #2892 (180e6f7) into master (b647fc5) will decrease coverage by 0.18%.
The diff coverage is 90.79%.

@@            Coverage Diff             @@
##           master    #2892      +/-   ##
==========================================
- Coverage   78.83%   78.64%   -0.19%     
==========================================
  Files         194      208      +14     
  Lines       26548    27212     +664     
==========================================
+ Hits        20928    21400     +472     
- Misses       5620     5812     +192

Flag	Coverage Δ
functest	`56.55% <89.45%> (+0.43%)`	⬆️
unittest	`69.18% <25.33%> (-1.49%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
wandb/sdk/internal/handler.py	`89.22% <66.66%> (-4.20%)`	⬇️
wandb/sdk/service/server.py	`82.22% <82.22%> (ø)`
wandb/sdk/service/service_base.py	`83.33% <83.33%> (ø)`
wandb/sdk/service/service_grpc.py	`84.09% <84.09%> (ø)`
wandb/sdk/interface/interface_shared.py	`84.83% <84.83%> (ø)`
wandb/sdk/interface/router.py	`96.72% <85.71%> (-1.59%)`	⬇️
wandb/sdk/service/service.py	`89.61% <90.00%> (-4.33%)`	⬇️
wandb/sdk/service/server_sock.py	`91.25% <91.25%> (ø)`
wandb/sdk/service/port_file.py	`92.45% <92.45%> (ø)`
wandb/sdk/internal/sender.py	`91.66% <93.33%> (-0.26%)`	⬇️
... and 38 more

kptkin · 2021-12-02T22:53:01Z

        parts = self._token_str.split("-")
-        assert len(parts) == 3, f"token must have 3 parts: {parts}"
+        assert len(parts) == 5, f"token must have 5 parts: {parts}"
        # TODO: make more robust?


what is this TODO mean?

TODO: what if there arent 5 parts, what if they are malformed with different versions and non supported tokens.
i think this is fine for now, but it is a note to care about this more in the future and think about how it should fail rather than an assert.. It "shouldnt" ever fail, but there are cases we already know if where it might.
Start your training job, before call wandb.init() someone updates the wandb release with pip install --upgrade.... Now the user process is speaking a potentially different dialect than the server process. Eventually we should actually handle these cases by bumping a version and making sure things are compatible.

kptkin · 2021-12-04T03:39:24Z

+            pass
+
+
+def _dict_from_pbmap(pbmap: "MessageMap[str, spb.SettingsValue]") -> Dict[str, Any]:


should it be under proto_util?

just being lazy, fixed in:
b99f463

kptkin · 2021-12-04T03:40:31Z

+        self.pid = 0
+
+    def run(self) -> None:
+        # TODO: catch exceptions and report errors to scheduler


do we track this TODO?

https://wandb.atlassian.net/browse/WB-7684
now we do

kptkin · 2021-12-04T03:45:44Z

+        self._target = target
+        self._kwargs = kwargs
+        self.daemon = True
+        self.pid = 0


what is this pid? why is it default to 0?

not needed right now. removed

kptkin · 2021-12-04T03:48:22Z

+if TYPE_CHECKING:
+    from google.protobuf.internal.containers import MessageMap
+
+    class GrpcServerType(object):


probably can remove this one, doesn't seem to be used

kptkin · 2021-12-04T03:50:03Z

+            relay_q=self._relay_q,
+            process=process,
+            process_check=False,
+            # use_router=False,


remove unused code?

kptkin · 2021-12-04T03:57:29Z

+
+        streams_to_join = []
+        while streams:
+            for sid, stream in list(streams.items()):


why do you need to wrap it in a list?

Added a comment:

# Note that we materialize the generator so we can modify the underlying list

kptkin · 2021-12-04T03:58:30Z

+        for stream in streams_to_join:
+            stream.join()
+
+        wandb.termlog("Done!")  # type: ignore


maybe more informative message?

This was covered in this ticket already, just havent taken the time to do it yet:
https://wandb.atlassian.net/browse/WB-6948

kptkin · 2021-12-04T04:00:12Z

+
+    def _loop(self) -> None:
+        while not self._stopped.is_set():
+            # TODO: check for parent process going away


Do you want to address this TODO in a different PR? should we track this?

Added to:
https://wandb.atlassian.net/browse/WB-7684

kptkin · 2021-12-04T04:01:24Z

+        self._action_q = queue.Queue()
+
+    def _get_stopped_event(self) -> "Event":
+        # TODO: clean this up, there should be a better way to abstract this


This refers to the Event logic we talked about? Do you want to track this one as well?

Added to:
https://wandb.atlassian.net/browse/WB-7684

kptkin · 2021-12-04T04:02:15Z

+- StreamRecord:
+- WandbServicer:
+
+"""


Probably want to update the docstring here?

Improved:

+"""streams: class that manages internal threads for each run. +StreamThread: Thread that runs internal.wandb_internal() +StreamRecord: All the external state for the internal thread (queues, etc) +StreamAction: Lightweight record for stream ops for thread safety with grpc +StreamMux: Container for dictionary of stream threads per runid """

kptkin · 2021-12-04T08:03:57Z

+
+    def server_record_publish(self, sreq: "spb.ServerRequest") -> None:
+        record = sreq.record_publish
+        # print("GOT rec", record)


remove commented out code

kptkin · 2021-12-04T08:04:07Z

+        iface.record_q.put(record)
+
+    def server_inform_finish(self, sreq: "spb.ServerRequest") -> None:
+        # print("serv INF FIN")


here as well :)

kptkin · 2021-12-04T08:04:14Z

+        self._mux.drop_stream(stream_id)
+
+    def server_inform_teardown(self, sreq: "spb.ServerRequest") -> None:
+        # print("serv INF TEARDOWN")


and this one

kptkin · 2021-12-04T08:04:45Z

+            # print("GOT", type(conn))
+            # print("Connected by", addr)


remove commented out code

kptkin · 2021-12-04T08:05:26Z

+        # print(f"Running at port: {self.port}")
+        self._thread = SockAcceptThread(sock=self._sock, mux=self._mux)
+        self._thread.start()
+        # self._dbg_thread = DebugThread(mux=self._mux)
+        # self._dbg_thread.start()


remove commented out code?

kptkin

Amazing!!! 🤩 🚀

raubitsj added 4 commits November 3, 2021 18:08

remame grpc_server

54968d7

Rename

d47b669

Large update -- needs cleanup

a94d5b0

Merge remote-tracking branch 'origin' into mp-add-fork-support

84edb41

raubitsj changed the title ~~[WIP] Add alternative service backend (not grpc) to support fork mode and higher performance~~ [WIP] Add service backend using sockets (support fork) Nov 7, 2021

raubitsj added 10 commits November 7, 2021 12:53

refactor service a bit

b07a6a1

Fix flake8 errors

33299d1

Fix many mypy errors

60db38a

more flake fixes

cf41245

More mypy checking

bb3e7e7

More annotation cleanups

95f20cf

Change mypy to be opt-out instead of opt-in

7f5438d

Merge remote-tracking branch 'origin/master' into mp-add-fork-support

c5d18c3

fix flake/mypy

6bf7dd7

add thread names and debug thread

9244617

raubitsj changed the title ~~[WIP] Add service backend using sockets (support fork)~~ [WIP][WB-6939] Add service backend using sockets (support fork) Nov 8, 2021

raubitsj added 13 commits November 8, 2021 05:06

Add grpc testing

d64b6e3

disable run status checker for now

d81bd62

temporarily allow manager

9f3dadc

Merge remote-tracking branch 'origin/master' into mp-add-fork-support

d92fb84

Merge remote-tracking branch 'origin/master' into mp-add-fork-support

7a9b196

Merge remote-tracking branch 'origin/master' into mp-add-fork-support

8a461df

reenable runchecker

c63cfea

Add some docs for wandb-service sock implementation

9cf2442

Remove debugging, improve shutdown

307ea0b

format

1b8b5a6

clean up lint

4344cf8

Merge remote-tracking branch 'origin/master' into mp-add-fork-support

b9a5a2d

fix shutdown

1d4b43b

raubitsj added 4 commits December 2, 2021 13:40

remove todo

8db3134

assert if ports not found

deb2c21

change debug default to True

f8da658

remove commented out code

cabdd50

kptkin reviewed Dec 2, 2021

View reviewed changes

Merge branch 'master' into mp-add-fork-support

a16c4d3

kptkin reviewed Dec 4, 2021

View reviewed changes

raubitsj added 5 commits December 3, 2021 21:19

Move settings_dict_from_pbmap to proto_util

b99f463

remove pid attribute (unused)

c157ce9

remove unused code

6b4aea1

update from pr comments

ce2bb9a

update docstring

9efc540

kptkin reviewed Dec 4, 2021

View reviewed changes

kptkin approved these changes Dec 4, 2021

View reviewed changes

final review fixes

180e6f7

raubitsj merged commit 31c12aa into master Dec 4, 2021

raubitsj deleted the mp-add-fork-support branch December 4, 2021 16:32

		pass


		def _dict_from_pbmap(pbmap: "MessageMap[str, spb.SettingsValue]") -> Dict[str, Any]:

Conversation

raubitsj commented Nov 7, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Testing

Uh oh!

codecov Bot commented Nov 7, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kptkin left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

raubitsj commented Nov 7, 2021 •

edited

Loading

codecov Bot commented Nov 7, 2021 •

edited

Loading