fix(sdk): add support for service running in a pex based environment (#…

…4440)

* fix pex

* python_executable -> settings

* add a test

* fix(sdk): pex pex pex

* fix config

* oomny v goroo ne poidyot

* oomny v goroo ne poidyot

* oora my lomim gnutsya shvedy

* oora my lomim gnutsya shvedy

* oora my lomim gnutsya shvedy

* oora my lomim gnutsya shvedy

* oora my lomim gnutsya shvedy

* oora my lomim gnutsya shvedy

* replace borky yea test with a (borky) circle job

* clean up

* lint

Co-authored-by: Dmitry Duev <dima@wandb.com>
Co-authored-by: Dmitry Duev <dmitryduev@users.noreply.github.com>

.circleci/config.yml

@@ -331,6 +331,38 @@ jobs:
                 # taken from slack-secrets context
                 channel: $SLACK_SDK_NIGHTLY_CI_CHANNEL
 
+  pex:
+    parameters:
+      python_version_major:
+        type: integer
+        default: 3
+      python_version_minor:
+        type: integer
+        default: 8
+    docker:
+      - image: "python:<<parameters.python_version_major>>.<<parameters.python_version_minor>>"
+    resource_class: small
+    working_directory: /mnt/ramdisk
+    steps:
+      - checkout
+      - run:
+          name: Install system deps
+          command: |
+            apt-get update && apt-get install -y libsndfile1 ffmpeg
+      - run:
+          name: Install pex
+          command: |
+            pip install -U pip
+            pip install pex
+          no_output_timeout: 5m
+      - run:
+          name: Run pex test
+          no_output_timeout: 5m
+          command: |
+            pex . -o wandb.pex
+            ./wandb.pex -c "import wandb; wandb.init(mode='offline'); wandb.finish()"
+      - save-test-results
+
   win:
     parameters:
       python_version_major:
@@ -1253,6 +1285,11 @@ workflows:
           toxenv: "unit-s_nb-py<<matrix.python_version_major>><<matrix.python_version_minor>>"
           parallelism: 1
           xdist: 1
+      #
+      # pex test
+      #
+      - pex:
+          name: "pex"
 
   # todo: needs love
   #  manual_test:

wandb/sdk/internal/settings_static.py

@@ -50,6 +50,7 @@ class SettingsStatic:
     anonymous: Optional[str]
     host: Optional[str]
     username: Optional[str]
+    _executable: str
 
     # TODO(jhr): clean this up, it is only in SettingsStatic and not in Settings
     _log_level: int

wandb/sdk/service/service.py

@@ -6,7 +6,6 @@
 import os
 import platform
 import subprocess
-import sys
 import tempfile
 import time
 from typing import Any, Dict, Optional
@@ -22,8 +21,13 @@ class _Service:
     _service_interface: ServiceInterface
     _internal_proc: Optional[subprocess.Popen]
 
-    def __init__(self, _use_grpc: bool = False) -> None:
+    def __init__(
+        self,
+        _python_executable: str,
+        _use_grpc: bool = False,
+    ) -> None:
         self._use_grpc = _use_grpc
+        self._python_executable = _python_executable
         self._stub = None
         self._grpc_port = None
         self._sock_port = None
@@ -82,7 +86,8 @@ def _launch_server(self) -> None:
             fname = os.path.join(tmpdir, f"port-{pid}.txt")
 
             pid_str = str(os.getpid())
-            exec_cmd_list = [sys.executable, "-m"]
+            executable = self._python_executable
+            exec_cmd_list = [executable, "-m"]
             # Add coverage collection if needed
             if os.environ.get("YEA_RUN_COVERAGE") and os.environ.get("COVERAGE_RCFILE"):
                 exec_cmd_list += ["coverage", "run", "-m"]

wandb/sdk/wandb_manager.py

@@ -101,7 +101,10 @@ def __init__(self, settings: "Settings", _use_grpc: bool = False) -> None:
         self._atexit_lambda = None
         self._hooks = None
 
-        self._service = service._Service(_use_grpc=_use_grpc)
+        self._service = service._Service(
+            _python_executable=settings._executable,
+            _use_grpc=_use_grpc,
+        )
 
         token = _ManagerToken.from_environment()
         if not token:

wandb/sdk/wandb_settings.py

@@ -1395,7 +1395,12 @@ def _infer_settings_from_environment(
                 # chroot jails or docker containers. Return user id in these cases.
                 settings["username"] = str(os.getuid())
 
-        settings["_executable"] = sys.executable
+        # one special case here is running inside a PEX environment,
+        # see https://pex.readthedocs.io/en/latest/index.html for more info about PEX
+        _executable = self._executable or os.environ.get("PEX") or sys.executable
+        if _executable is None or _executable == "":
+            _executable = "python3"
+        settings["_executable"] = _executable
 
         settings["docker"] = wandb.env.get_docker(wandb.util.image_id_from_k8s())