Repository: wandb/wandb
Pull request: fix(sdk): Prevent stalling during artifact download when WANDB_MODE=offline
Description
-----------
<!--
Include reference to internal ticket "Fixes WB-NNNNN" and/or GitHub issue "Fixes #NNNN" (if applicable)
-->
- Fixes #7960 
- Fixes WB-20039, WB-19978, WB-19883


This PR addresses an issue where the artifact download operation hangs indefinitely when using the new core backend with `WANDB_MODE` set to "offline". The solution involves updating the behavior to appropriately handle the offline mode, ensuring that the operation either proceeds correctly or fails with a clear error message.

### Fix
For this fix, we simply preempt the `download()` request that comes from the user (both from `wandb.Api().artifact("<name>").download()` and from the `wandb_run.py` functions (which internally use the `Api().artifact.download()` call anyway. 

I took reference from the `link_artifact` implementation, which currently raises a NotImplemented error when in `offline` mode [here](https://github.com/wandb/wandb/blob/main/wandb/sdk/wandb_run.py#L2940). That's not ideal, but the precedence is that we can preempt the request in user's Python process itself before it reaches on the backend, and throw a descriptive error. 

### Comparing Old vs New Behaviour

#### Old Behaviour
*Python Service*
- `WANDB_MODE` = "offline"
- `WANDB_REQUIRE_CORE` = "false"

Output:
```bash
wandb:   1 of 1 files downloaded.  
Artifact downloaded at: /tmp/wandb_artifacts/file.csv
```

Observation: The python backend incorrectly allows the user to download an artifact even when in offline mode.

*`core`*
- `WANDB_MODE` = "offline"
- `WANDB_REQUIRE_CORE` = "true"

Output:
```bash
wandb: Using wandb-core as the SDK backend. Please refer to https://wandb.me/wandb-core for more information.

```
Observation: Process stalls and does not exit

#### New Behaviour
*Python Service*
- `WANDB_MODE` = "offline"
- `WANDB_REQUIRE_CORE` = "false"

Output:
```bash
Traceback (most recent call last):
  File "/Users/anmol/work/repro/bug_7960/main.py", line 28, in <module>
    artifact_path = download_artifact(artifact_name)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/anmol/work/repro/bug_7960/main.py", line 14, in download_artifact
    artifact_dir = artifact.download(root="/tmp/wandb_artifacts")
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/anmol/work/projects/wandb/wandb/sdk/artifacts/artifact.py", line 1658, in download
    raise RuntimeError("Cannot download artifacts in offline mode.")
RuntimeError: Cannot download artifacts in offline mode.
```

#### `core`
- `WANDB_MODE` = "offline"
- `WANDB_REQUIRE_CORE` = "true"

Output:
```bash
Traceback (most recent call last):
  File "/Users/anmol/work/repro/bug_7960/main.py", line 28, in <module>
    artifact_path = download_artifact(artifact_name)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/anmol/work/repro/bug_7960/main.py", line 14, in download_artifact
    artifact_dir = artifact.download(root="/tmp/wandb_artifacts")
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/anmol/work/projects/wandb/wandb/sdk/artifacts/artifact.py", line 1658, in download
    raise RuntimeError("Cannot download artifacts in offline mode.")
RuntimeError: Cannot download artifacts in offline mode.
```

Observation: The user process now handles the request and checks if its in offline mode (both using `env` var and `util` data structure that contains command line args). The behaviour in both Python and `core` is the same. 

### Fixing a Segmentation Fault
This PR also contains a quick fix for a segmentation fault that occurs in `sendRequestDownloadArtifact` method in `core`. This segfault is not triggered in the normal workflow, but only when:

- A download() request makes it through the user process to `core` unchecked
- Which leads to `graphqlClient` to be passed as `<nil>`, which in turn leads to a segfault

This achieves two goals:
- Prevents a segfault
- Adds a redundant check in the Go `core` if the request was somehow made in `offline` mode. 

This segfault fix is inspired by the precedence set in `sendREquestStopStatus` function [here](https://github.com/wandb/wandb/blob/main/core/pkg/server/sender.go#L1359). 

P.S: I do not think that this is an optimal way to handle these segfaults. Ideally the requests coming to these functions must be validated outside these functions, and broken/problematic requests should not make it to these functions at all. Technically, the changes in `sendRequestDownloadArtifact` will never be hit because we won't let the request come to this function, and throw an error to the user from the Python process itself. We should ideally do the same with `sendRequestStopStatus`. 

### Why does it stall?
As for the actual "stall" that occurs when trying to handle this request, I wasn't able to debug it properly. The issue arises in the `writer.go` file [here](https://github.com/wandb/wandb/blob/main/core/pkg/server/writer.go#L173). After it returns, the process does not return back anything back to the Python handler, and lets it stall forever. 

This behaviour was reproduced with both `download()` and `link_artifact()` requests, if we remove any validation from the Python processes, and let the problematic request go through into `core`. 

<!--
NEW: We're using a new changelog format that's more useful for users. Please
see CHANGELOG.md for details and update on relevant changes such as feature
additions, bug fixes, or removals/deprecations.
-->
- [x] I updated CHANGELOG.md, or it's not applicable


Testing
-------
How was this PR tested?
- Verified that the `artifact.download()` method raises an appropriate error when `WANDB_MODE` is set to "offline".
- Ensured that the method works correctly when `WANDB_MODE` is not set to "offline".
- Ran the existing test suite to confirm no regressions were introduced.

Minimal code to reproduce:

```python
import os

import wandb
wandb.require("core")
os.environ["WANDB_MODE"] = "offline"
os.environ["WANDB_CORE_DEBUG"] = "true"
# Set WANDB environment variables


# Download artifact function
def download_artifact(artifact_name):
    artifact = wandb.Api().artifact(artifact_name)
    artifact_dir = artifact.download(root="/tmp/wandb_artifacts")
    artifact_path = os.path.join(artifact_dir, "file.csv")
    return artifact_path


if __name__ == "__main__":
    # Replace 'your_artifact_name' with the actual artifact name
    artifact_name = "reg-team-2/collection-linking-quickstart/embeddings-test-dataset:v4"

    # let's try restore as well using a run
    run = wandb.init(project="collection-linking-quickstart")
    run.restore(name="output.log", replace=True, run_path="reg-team-2/collection-linking-quickstart/runs/zww2mukd")

    # Download the artifact
    artifact_path = download_artifact(artifact_name)
    print(f"Artifact downloaded at: {artifact_path}")
```
<!--
Ensure PR title compliance with the [conventional commits standards](https://github.com/wandb/wandb/blob/main/CONTRIBUTING.md#conventional-commits)
-->


[WB-20039]: https://wandb.atlassian.net/browse/WB-20039?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ

[WB-19978]: https://wandb.atlassian.net/browse/WB-19978?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ
[WB-19883]: https://wandb.atlassian.net/browse/WB-19883?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ

Problem statement:
Artifact download stalls with wandb core in offline mode
### Describe the bug

When I use wandb core and download an artifact in `offline` mode, the download hangs forever.

```python
import os

import wandb


# Download artifact function
def download_artifact(artifact_name):
    artifact = wandb.Api().artifact(artifact_name)
    artifact_dir = artifact.download(root="/tmp/wandb_artifacts")
    artifact_path = os.path.join(artifact_dir, "model.ckpt")
    return artifact_path


if __name__ == "__main__":
    wandb.require("core")

    # Set WANDB environment variables
    os.environ["WANDB_MODE"] = "offline"

    # Replace 'your_artifact_name' with the actual artifact name
    artifact_name = ""

    # Download the artifact
    artifact_path = download_artifact(artifact_name)
    print(f"Artifact downloaded at: {artifact_path}")
```
Commenting out `wandb.require("core")` or `os.environ["WANDB_MODE"] = "offline"` resolves the issue

Running strace on the script above shows the following loop:
```
futex(0x27df110, FUTEX_WAIT_BITSET_PRIVATE, 0, {tv_sec=1113302, tv_nsec=674199828}, FUTEX_BITSET_MATCH_ANY) = -1 ETIMEDOUT (Connection timed out)
futex(0x27df110, FUTEX_WAIT_BITSET_PRIVATE, 0, {tv_sec=1113303, tv_nsec=674394303}, FUTEX_BITSET_MATCH_ANY) = -1 ETIMEDOUT (Connection timed out)
futex(0x27df110, FUTEX_WAIT_BITSET_PRIVATE, 0, {tv_sec=1113304, tv_nsec=674627009}, FUTEX_BITSET_MATCH_ANY) = -1 ETIMEDOUT (Connection timed out)
futex(0x27df110, FUTEX_WAIT_BITSET_PRIVATE, 0, {tv_sec=1113305, tv_nsec=674878871}, FUTEX_BITSET_MATCH_ANY) = -1 ETIMEDOUT (Connection timed out)
futex(0x27df110, FUTEX_WAIT_BITSET_PRIVATE, 0, {tv_sec=1113306, tv_nsec=675080555}, FUTEX_BITSET_MATCH_ANY) = -1 ETIMEDOUT (Connection timed out)
poll([{fd=5, events=POLLOUT}], 1, 1000) = 1 ([{fd=5, revents=POLLOUT}])
sendto(5, "W\32\0\0\0\n\30\202\1\2\20\1\242\6\3\212\1\0\302\f\n\n\10kdf7290v", 31, 0, NULL, 0) = 31
```


### Additional Files

_No response_

### Environment

WandB version: wandb, version 0.17.4

OS: ubuntu:22.04

Python version: 3.11

Versions of relevant libraries:


### Additional Context

_No response_

Hints from issue discussion:
cc @dmitryduev @timoffex 
This is a regression from the non-core wandb backend, so it's likely a blocker to version 18 where the core backend will become the default.
@jackdent thanks for reporting the issue, we escalated it internally and will have one of the engineers look into it. 
Will keep you updated as soon as this issue is resolved. Thanks a lot of trying `wandb-core` and we should have a resolution for you very soon.

Commit message: removed problematic offline check with the settings structure mimicing run implementation

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f8b3844f46a..8ae212ec9db 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -34,6 +34,7 @@ Please add to the relevant subsections under Unreleased below on every PR where
 
 ### Fixed
 
+- Fix process stalling when trying to call `artifact.download()` when `WANDB_MODE = offline` in https://github.com/wandb/wandb/pull/8009
 - Handle `path_prefix`es that don't correspond to directory names when downloading artifacts by @moredatarequired in https://github.com/wandb/wandb/pull/7721
 - Fix creating or updating an artifact with the `incremental=True` flag by @amusipatla-wandb in https://github.com/wandb/wandb/pull/7939
 - Use filled resource_arg macros when submitting W&B Launch jobs to AmazonSageMaker by @KyleGoyette https://github.com/wandb/wandb/pull/7993
diff --git a/core/pkg/server/sender.go b/core/pkg/server/sender.go
index 98d36a83555..9a2481bb0c6 100644
--- a/core/pkg/server/sender.go
+++ b/core/pkg/server/sender.go
@@ -1265,7 +1265,13 @@ func (s *Sender) sendRequestLogArtifact(record *service.Record, msg *service.Log
 func (s *Sender) sendRequestDownloadArtifact(record *service.Record, msg *service.DownloadArtifactRequest) {
 	var response service.DownloadArtifactResponse
 
-	if err := artifacts.NewArtifactDownloader(
+	// Perform the nil check on the GraphQL client.
+	if s.graphqlClient == nil {
+		// Offline mode handling:
+		s.logger.Error("sender: sendRequestDownloadArtifact: cannot download artifact in offline mode")
+		response.ErrorMessage = "Artifact downloads are not supported in offline mode. " +
+			"Run 'wandb sync' after the run completes to sync the request."
+	} else if err := artifacts.NewArtifactDownloader(
 		s.ctx,
 		s.graphqlClient,
 		s.fileTransferManager,
@@ -1275,11 +1281,13 @@ func (s *Sender) sendRequestDownloadArtifact(record *service.Record, msg *servic
 		msg.SkipCache,
 		msg.PathPrefix,
 	).Download(); err != nil {
+		// Online mode handling: error during download
 		s.logger.CaptureError(
 			fmt.Errorf("sender: failed to download artifact: %v", err))
 		response.ErrorMessage = err.Error()
 	}
 
+	// Send the response using s.respond().
 	s.respond(record,
 		&service.Response{
 			ResponseType: &service.Response_DownloadArtifactResponse{
diff --git a/tests/pytest_tests/system_tests/test_artifacts/test_wandb_artifacts_api.py b/tests/pytest_tests/system_tests/test_artifacts/test_wandb_artifacts_api.py
index 7dea2d06e05..c47d63afc09 100644
--- a/tests/pytest_tests/system_tests/test_artifacts/test_wandb_artifacts_api.py
+++ b/tests/pytest_tests/system_tests/test_artifacts/test_wandb_artifacts_api.py
@@ -1,4 +1,6 @@
 import os
+import tempfile
+from unittest.mock import patch
 
 import pytest
 import wandb
@@ -36,6 +38,33 @@ def test_fetching_artifact_files(user, wandb_init):
     assert open(file_path).read() == "testing"
 
 
+def test_artifact_download_offline_mode(user, wandb_init, monkeypatch):
+    project = "test"
+
+    with tempfile.TemporaryDirectory() as tempdir:
+        # Create the test file in the temporary directory
+        filename = os.path.join(tempdir, "boom.txt")
+        with open(filename, "w") as f:
+            f.write("testing")
+
+        # Use monkeypatch to set WANDB_MODE
+        monkeypatch.setenv("WANDB_MODE", "offline")
+
+        # Mock the wandb.Api().artifact call to prevent actual network requests
+        with patch("wandb.Api.artifact") as mock_artifact:
+            mock_artifact.return_value = wandb.Artifact("test-artifact", "test-type")
+
+            with wandb_init(entity=user, project=project) as run:
+                artifact = wandb.Artifact("test-artifact", "test-type")
+                artifact.add_file(filename, "test-name")
+                run.log_artifact(artifact, aliases=["sequence"])
+
+            with pytest.raises(
+                RuntimeError, match="Cannot download artifacts in offline mode."
+            ):
+                artifact.download()
+
+
 def test_save_aliases_after_logging_artifact(user, wandb_init):
     project = "test"
     run = wandb_init(entity=user, project=project)
diff --git a/wandb/env.py b/wandb/env.py
index 702b3fd8e5a..35eda71ccb0 100644
--- a/wandb/env.py
+++ b/wandb/env.py
@@ -161,6 +161,12 @@ def is_debug(default: Optional[str] = None, env: Optional[Env] = None) -> bool:
     return _env_as_bool(DEBUG, default=default, env=env)
 
 
+def is_offline(env: Optional[Env] = None) -> bool:
+    if env is None:
+        env = os.environ
+    return env.get(MODE) == "offline"
+
+
 def error_reporting_enabled() -> bool:
     return _env_as_bool(ERROR_REPORTING, default="True")
 
diff --git a/wandb/sdk/artifacts/artifact.py b/wandb/sdk/artifacts/artifact.py
index 9b05b9f8c6b..69df4ebce45 100644
--- a/wandb/sdk/artifacts/artifact.py
+++ b/wandb/sdk/artifacts/artifact.py
@@ -75,6 +75,7 @@
 from wandb.sdk.lib.mailbox import Mailbox
 from wandb.sdk.lib.paths import FilePathStr, LogicalPath, StrPath, URIStr
 from wandb.sdk.lib.runid import generate_id
+from wandb.sdk.wandb_settings import Settings
 from wandb.util import get_core_path
 
 reset_path = util.vendor_setup()
@@ -191,6 +192,7 @@ def __init__(
         self._created_at: Optional[str] = None
         self._updated_at: Optional[str] = None
         self._final: bool = False
+        self._settings: Settings
 
         # Cache.
         artifact_instance_cache[self._client_id] = self
@@ -1647,12 +1649,16 @@ def download(
 
         Raises:
             ArtifactNotLoggedError: If the artifact is not logged.
+            RuntimeError: If the artifact is attempted to be downloaded in offline mode.
         """
         self._ensure_logged("download")
 
         root = FilePathStr(str(root or self._default_root()))
         self._add_download_root(root)
 
+        if env.is_offline() or self._settings._offline():
+            raise RuntimeError("Cannot download artifacts in offline mode.")
+
         if is_require_core():
             return self._download_using_core(
                 root=root,


Code review comments:
diff_hunk: @@ -1647,12 +1649,16 @@ def download(
 
         Raises:
             ArtifactNotLoggedError: If the artifact is not logged.
+            RuntimeError: If the artifact is attempted to be downloaded in offline mode.
         """
         self._ensure_logged("download")
 
         root = FilePathStr(str(root or self._default_root()))
         self._add_download_root(root)
 
+        if env.is_offline() or self._settings._offline():
path: wandb/sdk/artifacts/artifact.py
text: @user1:
Instead of assigning to a `_settings` member, I think you can look at the `wandb.settings` value here, similar to how the existence of a run is checked when logging

@author:
Can you suggest the fix? I am having a hard time pinning down the exact ways to check for offline. The codebase is quite inconsistent or not nearly as proactive in doing so. 

@author:
After discussion we're settling on `util._is_offline()`

diff_hunk: @@ -191,6 +192,7 @@ def __init__(
         self._created_at: Optional[str] = None
         self._updated_at: Optional[str] = None
         self._final: bool = False
+        self._settings: Settings
path: wandb/sdk/artifacts/artifact.py
text: @user1:
Does this work? I think the `Settings` object is a singleton, but it's still not instantiated here.