Repository: wandb/wandb
Pull request: fix(sdk): unify offline and online mode during init and fix multiprocess attach
…xisting run (#5144)

Fixes WB-NNNN
Fixes #5071

Description
-----------
What does the PR do?

TODO:
- [x] unify `deliver_run` and `publish_run`
- [x] move attach handling from sender to handler
- [x] unify offline_run and online_run in the run class

Testing
-------
How was this PR tested?

Checklist
-------
- [ ] Include reference to internal ticket "Fixes WB-NNNN" and/or GitHub issue "Fixes #NNNN" (if applicable)
- [ ] Ensure PR title compliance with the [conventional commits standards](https://github.com/wandb/wandb/blob/main/CONTRIBUTING.md#conventional-commits)


Problem statement:
[CLI]: wandb.errors.UsageError: Timeout attaching to run when in offline mode
### Describe the bug

<!--- Description of the issue below  -->
When I try to access the config of the run object, I get the exception below. However, this only happens when in offline mode; it works fine when in online mode. (note that I am using multiprocessing as shown below and cannot seem to reproduce in the simple two-line code of a single process run=wandb.init; config=run.config). It only works when in online mode, and I can repro 100% when in offline mode.

<!--- A minimal code snippet between the quotes below  -->
```python
import multiprocessing
import wandb

def simple_bug_f(wandb_run):
    print("Trying to get config")
    config = wandb_run.config
    print("Got config")

if __name__ == '__main__':
    wandb_run = wandb.init(mode="offline")
    
    ctx = multiprocessing.get_context('spawn')
    pool = ctx.Pool(processes=2)

    futures = []
    for i in range(2):
        future = pool.apply_async(simple_bug_f, args=(wandb_run,))
        futures.append(future)
    
    [future.get() for future in futures]

```

<!--- A full traceback of the exception in the quotes below -->
```shell
    config = wandb_run.config
  File "python3.8/site-packages/wandb/sdk/wandb_run.py", line 332, in wrapper
    raise e
  File "python3.8/site-packages/wandb/sdk/wandb_run.py", line 327, in wrapper
    wandb._attach(run=self)
  File "python3.8/site-packages/wandb/sdk/wandb_init.py", line 897, in _attach
    raise UsageError("Timeout attaching to run")
wandb.errors.UsageError: Timeout attaching to run
```


### Additional Files

_No response_

### Environment

WandB version: 0.13.10

OS: Linux

Python version: 3.8

Versions of relevant libraries:

**Also tried in another totally different environment which has python 3.9.10** I see the problem exactly the same

### Additional Context

_No response_

Hints from issue discussion:
Did some digging and I can see that _attach in wandb_init.py does not have any offline handling code and always assumes online. As per a TODO comment, this code needs to be consolidated with the wandb.init() path which has handling for offline mode. 
@RamyE Thanks for reporting this issue! I will take a look at this, but your observation seems to be correct.
Hey @kptkin,
Can I pick this up ? 
Hey @techytushar sure you can! Thank you very much for proposing your help!
Bill Morrisson commented: 
Hello @techytushar, have you been able to take a look at it?

@billmdevs I didn't get a chance to look into this, will be working on this over the weekend 
Alright @techytushar  thank you very much!!! I see you pushed a PR. A member of our team is going to review and/or merge or get back to you with some queries.

Commit message: fix typing

Patch to review:
diff --git a/tests/functional_tests/t0_main/offline/t2_multiprocess.py b/tests/functional_tests/t0_main/offline/t2_multiprocess.py
new file mode 100644
index 00000000000..df292f0a19c
--- /dev/null
+++ b/tests/functional_tests/t0_main/offline/t2_multiprocess.py
@@ -0,0 +1,39 @@
+#!/usr/bin/env python
+"""Offline runs running in different processes.
+
+---
+id: 0.offline.02-multiprocess
+env:
+  - WANDB_BASE_URL: https://does.not-resolve/
+tag:
+  skip: true
+command:
+  timeout: 20
+plugin:
+  - wandb
+assert:
+  - :wandb:runs_len: 1
+  - :yea:exit: 0
+"""
+
+import multiprocessing
+
+import wandb
+
+
+def log_metric(run):
+    run.log({"metric": 10})
+
+
+if __name__ == "__main__":
+    run = wandb.init(mode="offline")
+
+    ctx = multiprocessing.get_context("spawn")
+    pool = ctx.Pool(processes=2)
+
+    futures = []
+    for _ in range(2):
+        future = pool.apply_async(log_metric, args=(run,))
+        futures.append(future)
+
+    [future.get() for future in futures]
diff --git a/wandb/sdk/internal/handler.py b/wandb/sdk/internal/handler.py
index 81c9827ecdc..61a318d9cf4 100644
--- a/wandb/sdk/internal/handler.py
+++ b/wandb/sdk/internal/handler.py
@@ -191,6 +191,11 @@ def handle_request_login(self, record: Record) -> None:
         self._dispatch_record(record)
 
     def handle_run(self, record: Record) -> None:
+        if self._settings._offline:
+            self._run_proto = record.run
+            result = proto_util._result_from_record(record)
+            result.run_result.run.CopyFrom(record.run)
+            self._respond_result(result)
         self._dispatch_record(record)
 
     def handle_stats(self, record: Record) -> None:
@@ -627,7 +632,12 @@ def handle_request_check_version(self, record: Record) -> None:
         self._dispatch_record(record)
 
     def handle_request_attach(self, record: Record) -> None:
-        self._dispatch_record(record)
+        result = proto_util._result_from_record(record)
+        attach_id = record.request.attach.attach_id
+        assert attach_id
+        assert self._run_proto
+        result.response.attach_response.run.CopyFrom(self._run_proto)
+        self._respond_result(result)
 
     def handle_request_log_artifact(self, record: Record) -> None:
         self._dispatch_record(record)
@@ -678,6 +688,8 @@ def handle_request_run_start(self, record: Record) -> None:
         assert run_start
         assert run_start.run
 
+        self._run_proto = run_start.run
+
         self._run_start_time = run_start.run.start_time.ToMicroseconds() / 1e6
 
         self._track_time = time.time()
diff --git a/wandb/sdk/internal/sender.py b/wandb/sdk/internal/sender.py
index 70465f99dd2..a5f7aa15b33 100644
--- a/wandb/sdk/internal/sender.py
+++ b/wandb/sdk/internal/sender.py
@@ -486,24 +486,6 @@ def send_request_check_version(self, record: "Record") -> None:
                 result.response.check_version_response.delete_message = delete_message
         self._respond_result(result)
 
-    def _send_request_attach(
-        self,
-        req: wandb_internal_pb2.AttachRequest,
-        resp: wandb_internal_pb2.AttachResponse,
-    ) -> None:
-        attach_id = req.attach_id
-        assert attach_id
-        assert self._run
-        resp.run.CopyFrom(self._run)
-
-    def send_request_attach(self, record: "Record") -> None:
-        assert record.control.req_resp or record.control.mailbox_slot
-        result = proto_util._result_from_record(record)
-        self._send_request_attach(
-            record.request.attach, result.response.attach_response
-        )
-        self._respond_result(result)
-
     def send_request_stop_status(self, record: "Record") -> None:
         result = proto_util._result_from_record(record)
         status_resp = result.response.stop_status_response
diff --git a/wandb/sdk/wandb_init.py b/wandb/sdk/wandb_init.py
index d8ccb3cbdf5..13267fbac7d 100644
--- a/wandb/sdk/wandb_init.py
+++ b/wandb/sdk/wandb_init.py
@@ -717,79 +717,79 @@ def init(self) -> Union[Run, RunDisabled, None]:  # noqa: C901
             with telemetry.context(run=run) as tel:
                 tel.feature.offline = True
 
-            backend.interface.publish_run(run_proto)
-            run._set_run_obj_offline(run_proto)
+            # backend.interface.publish_run(run_proto)
+            # run._set_run_obj_offline(run_proto)
             if self.settings.resume:
                 wandb.termwarn(
                     "`resume` will be ignored since W&B syncing is set to `offline`. "
                     f"Starting a new run with run id {run.id}."
                 )
-        else:
-            error: Optional["wandb.errors.Error"] = None
+        # else:
+        error: Optional["wandb.errors.Error"] = None
+
+        timeout = self.settings.init_timeout
 
-            timeout = self.settings.init_timeout
+        logger.info(f"communicating run to backend with {timeout} second timeout")
 
-            logger.info(f"communicating run to backend with {timeout} second timeout")
+        run_init_handle = backend.interface.deliver_run(run_proto)
+        result = run_init_handle.wait(
+            timeout=timeout,
+            on_progress=self._on_progress_init,
+            cancel=True,
+        )
+        if result:
+            run_result = result.run_result
 
-            run_init_handle = backend.interface.deliver_run(run_proto)
-            result = run_init_handle.wait(
-                timeout=timeout,
-                on_progress=self._on_progress_init,
-                cancel=True,
+        if run_result is None:
+            error_message = (
+                f"Run initialization has timed out after {timeout} sec. "
+                f"\nPlease refer to the documentation for additional information: {wburls.get('doc_start_err')}"
+            )
+            # We're not certain whether the error we encountered is due to an issue
+            # with the server (a "CommError") or if it's a problem within the SDK (an "Error").
+            # This means that the error could be a result of the server being unresponsive,
+            # or it could be because we were unable to communicate with the wandb service.
+            error = CommError(error_message)
+            run_init_handle._cancel()
+        elif run_result.HasField("error"):
+            error = ProtobufErrorHandler.to_exception(run_result.error)
+
+        if error is not None:
+            logger.error(f"encountered error: {error}")
+            if not manager:
+                # Shutdown the backend and get rid of the logger
+                # we don't need to do console cleanup at this point
+                backend.cleanup()
+                self.teardown()
+            raise error
+
+        assert run_result is not None  # for mypy
+
+        if not run_result.HasField("run"):
+            raise Error(
+                "It appears that something have gone wrong during the program execution as an unexpected missing field was encountered. "
+                "(run_result is missing the 'run' field)"
             )
-            if result:
-                run_result = result.run_result
 
-            if run_result is None:
-                error_message = (
-                    f"Run initialization has timed out after {timeout} sec. "
-                    f"\nPlease refer to the documentation for additional information: {wburls.get('doc_start_err')}"
-                )
-                # We're not certain whether the error we encountered is due to an issue
-                # with the server (a "CommError") or if it's a problem within the SDK (an "Error").
-                # This means that the error could be a result of the server being unresponsive,
-                # or it could be because we were unable to communicate with the wandb service.
-                error = CommError(error_message)
-                run_init_handle._cancel()
-            elif run_result.HasField("error"):
-                error = ProtobufErrorHandler.to_exception(run_result.error)
-
-            if error is not None:
-                logger.error(f"encountered error: {error}")
-                if not manager:
-                    # Shutdown the backend and get rid of the logger
-                    # we don't need to do console cleanup at this point
-                    backend.cleanup()
-                    self.teardown()
-                raise error
-
-            assert run_result is not None  # for mypy
-
-            if not run_result.HasField("run"):
-                raise Error(
-                    "It appears that something have gone wrong during the program execution as an unexpected missing field was encountered. "
-                    "(run_result is missing the 'run' field)"
-                )
+        if run_result.run.resumed:
+            logger.info("run resumed")
+            with telemetry.context(run=run) as tel:
+                tel.feature.resumed = run_result.run.resumed
 
-            if run_result.run.resumed:
-                logger.info("run resumed")
-                with telemetry.context(run=run) as tel:
-                    tel.feature.resumed = run_result.run.resumed
+        run._set_run_obj(run_result.run)
 
-            run._set_run_obj(run_result.run)
-            run._on_init()
+        run._on_init()
 
         logger.info("starting run threads in backend")
         # initiate run (stats and metadata probing)
-        run_obj = run._run_obj or run._run_obj_offline
 
         if manager:
             manager._inform_start(settings=self.settings, run_id=self.settings.run_id)
 
         assert backend.interface
-        assert run_obj
+        assert run._run_obj
 
-        run_start_handle = backend.interface.deliver_run_start(run_obj)
+        run_start_handle = backend.interface.deliver_run_start(run._run_obj)
         # TODO: add progress to let user know we are doing something
         run_start_result = run_start_handle.wait(timeout=30)
         if run_start_result is None:
@@ -872,11 +872,7 @@ def _attach(
 
     settings: Settings = copy.copy(_wl._settings)
     settings.update(
-        {
-            "run_id": attach_id,
-            "_start_time": response["_start_time"],
-            "_start_datetime": response["_start_datetime"],
-        },
+        {"run_id": attach_id, **response},
         source=Source.INIT,
     )
 
@@ -896,6 +892,7 @@ def _attach(
     assert backend.interface
 
     mailbox.enable_keepalive()
+
     attach_handle = backend.interface.deliver_attach(attach_id)
     # TODO: add progress to let user know we are doing something
     attach_result = attach_handle.wait(timeout=30)
@@ -905,6 +902,7 @@ def _attach(
     attach_response = attach_result.response.attach_response
     if attach_response.error and attach_response.error.message:
         raise UsageError(f"Failed to attach to run: {attach_response.error.message}")
+
     run._set_run_obj(attach_response.run)
     run._on_attach()
     return run
diff --git a/wandb/sdk/wandb_run.py b/wandb/sdk/wandb_run.py
index 47343b8dc5c..7caacafd187 100644
--- a/wandb/sdk/wandb_run.py
+++ b/wandb/sdk/wandb_run.py
@@ -480,7 +480,6 @@ class Run:
     _notes: Optional[str]
 
     _run_obj: Optional[RunRecord]
-    _run_obj_offline: Optional[RunRecord]
     # Use string literal annotation because of type reference loop
     _backend: Optional["wandb.sdk.backend.backend.Backend"]
     _internal_run_interface: Optional[
@@ -603,7 +602,6 @@ def _init(
 
         # Returned from backend request_run(), set from wandb_init?
         self._run_obj = None
-        self._run_obj_offline = None
 
         # Created when the run "starts".
         self._run_status_checker = None
@@ -913,9 +911,8 @@ def tags(self) -> Optional[Tuple]:
         """Tags associated with the run, if there are any."""
         if self._tags:
             return self._tags
-        run_obj = self._run_obj or self._run_obj_offline
-        if run_obj:
-            return tuple(run_obj.tags)
+        if self._run_obj:
+            return tuple(self._run_obj.tags)
         return None
 
     @tags.setter
@@ -998,8 +995,7 @@ def step(self) -> int:
         return self._step
 
     def project_name(self) -> str:
-        run_obj = self._run_obj or self._run_obj_offline
-        return run_obj.project if run_obj else ""
+        return self._run_obj.project if self._run_obj else ""
 
     @property
     @_run_decorator._attach
@@ -1025,8 +1021,7 @@ def disabled(self) -> bool:
         return self._settings._noop
 
     def _get_group(self) -> str:
-        run_obj = self._run_obj or self._run_obj_offline
-        return run_obj.run_group if run_obj else ""
+        return self._run_obj.run_group if self._run_obj else ""
 
     @property
     @_run_decorator._attach
@@ -1037,7 +1032,7 @@ def group(self) -> str:
 
         If you are doing a distributed training you should give all of the
             runs in the training the same group.
-        If you are doing crossvalidation you should give all the crossvalidation
+        If you are doing cross-validation you should give all the cross-validation
             folds the same group.
         """
         return self._get_group()
@@ -1045,8 +1040,7 @@ def group(self) -> str:
     @property
     @_run_decorator._attach
     def job_type(self) -> str:
-        run_obj = self._run_obj or self._run_obj_offline
-        return run_obj.job_type if run_obj else ""
+        return self._run_obj.job_type if self._run_obj else ""
 
     @property
     @_run_decorator._attach
@@ -1456,6 +1450,9 @@ def _set_teardown_hooks(self, hooks: List[TeardownHook]) -> None:
 
     def _set_run_obj(self, run_obj: RunRecord) -> None:
         self._run_obj = run_obj
+        if self.settings._offline:
+            return
+
         self._entity = run_obj.entity
         self._project = run_obj.project
 
@@ -1484,9 +1481,6 @@ def _set_run_obj(self, run_obj: RunRecord) -> None:
             settings=self._settings,
         )
 
-    def _set_run_obj_offline(self, run_obj: RunRecord) -> None:
-        self._run_obj_offline = run_obj
-
     def _add_singleton(
         self, data_type: str, key: str, value: Dict[Union[int, str], str]
     ) -> None:
@@ -2172,6 +2166,8 @@ def _console_stop(self) -> None:
             self._output_writer = None
 
     def _on_init(self) -> None:
+        if self._settings._offline:
+            return
         if self._backend and self._backend.interface:
             logger.info("communicating current version")
             version_handle = self._backend.interface.deliver_check_version(
@@ -2945,10 +2941,9 @@ def _log_artifact(
 
     def _public_api(self, overrides: Optional[Dict[str, str]] = None) -> PublicApi:
         overrides = {"run": self._run_id}
-        run_obj = self._run_obj
-        if run_obj is not None:
-            overrides["entity"] = run_obj.entity
-            overrides["project"] = run_obj.project
+        if not (self._settings._offline or self._run_obj is None):
+            overrides["entity"] = self._run_obj.entity
+            overrides["project"] = self._run_obj.project
         return public.Api(overrides)
 
     # TODO(jhr): annotate this


Code review comments:
diff_hunk: @@ -717,79 +717,79 @@ def init(self) -> Union[Run, RunDisabled, None]:  # noqa: C901
             with telemetry.context(run=run) as tel:
                 tel.feature.offline = True
 
-            backend.interface.publish_run(run_proto)
-            run._set_run_obj_offline(run_proto)
+            # backend.interface.publish_run(run_proto)
+            # run._set_run_obj_offline(run_proto)
             if self.settings.resume:
                 wandb.termwarn(
                     "`resume` will be ignored since W&B syncing is set to `offline`. "
                     f"Starting a new run with run id {run.id}."
                 )
-        else:
-            error: Optional["wandb.errors.Error"] = None
+        # else:
path: wandb/sdk/wandb_init.py
text: @author:
check offline + resume with the new logic

diff_hunk: @@ -717,79 +717,79 @@ def init(self) -> Union[Run, RunDisabled, None]:  # noqa: C901
             with telemetry.context(run=run) as tel:
                 tel.feature.offline = True
 
-            backend.interface.publish_run(run_proto)
-            run._set_run_obj_offline(run_proto)
+            # backend.interface.publish_run(run_proto)
+            # run._set_run_obj_offline(run_proto)
             if self.settings.resume:
                 wandb.termwarn(
                     "`resume` will be ignored since W&B syncing is set to `offline`. "
                     f"Starting a new run with run id {run.id}."
                 )
-        else:
-            error: Optional["wandb.errors.Error"] = None
+        # else:
+        error: Optional["wandb.errors.Error"] = None
+
+        timeout = self.settings.init_timeout
 
-            timeout = self.settings.init_timeout
+        logger.info(f"communicating run to backend with {timeout} second timeout")
 
-            logger.info(f"communicating run to backend with {timeout} second timeout")
+        run_init_handle = backend.interface.deliver_run(run_proto)
path: wandb/sdk/wandb_init.py
text: @author:
undo the `run_proto` input change