xoscar 0.6.1__cp310-cp310-macosx_10_9_x86_64.whl → 0.6.2__cp310-cp310-macosx_10_9_x86_64.whl

xoscar/backends/indigen/pool.py

@@ -134,6 +134,31 @@ class SubpoolStatus:
     traceback: TracebackType | None = None
 
 
+_PRE_SET_ENV_LOCK = asyncio.Lock()
+
+
+@contextlib.asynccontextmanager
+async def _pre_set_env_in_main(env: dict[str, str]):
+    # Normally, `env` is set in sub pool,
+    # but something may have happened during initialization,
+    # e.g. CUDA_VISIBLE_DEVICES is too late to set when some actions may have polluted cuda
+    # we have to set environ before new process started
+    # enable this only when XOSCAR_PRE_SET_ENV=1
+    enable_pre_set_env = bool(int(os.getenv("XOSCAR_PRE_SET_ENV", 0)))
+    if not enable_pre_set_env or not env:
+        yield
+        return
+
+    global_environ = os.environ.copy()
+    async with _PRE_SET_ENV_LOCK:
+        try:
+            logger.debug("Updating environment variables in main: %s", env)
+            os.environ.update(env)
+            yield
+        finally:
+            os.environ = global_environ  # type: ignore
+
+
 @_register_message_handler
 class MainActorPool(MainActorPoolBase):
     @classmethod
@@ -223,9 +248,12 @@ class MainActorPool(MainActorPoolBase):
 
         _patch_spawn_get_preparation_data()
         loop = asyncio.get_running_loop()
-        with futures.ThreadPoolExecutor(1) as executor:
-            create_pool_task = loop.run_in_executor(executor, start_pool_in_process)
-            return await create_pool_task
+        async with _pre_set_env_in_main(
+            actor_pool_config.get_pool_config(process_index)["env"]
+        ):
+            with futures.ThreadPoolExecutor(1) as executor:
+                create_pool_task = loop.run_in_executor(executor, start_pool_in_process)
+                return await create_pool_task
 
     @classmethod
     async def wait_sub_pools_ready(cls, create_pool_tasks: List[asyncio.Task]):
@@ -394,9 +422,10 @@ class MainActorPool(MainActorPoolBase):
             return process, process_status
 
         loop = asyncio.get_running_loop()
-        with futures.ThreadPoolExecutor(1) as executor:
-            create_pool_task = loop.run_in_executor(executor, start_pool_in_process)
-            process, process_status = await create_pool_task
+        async with _pre_set_env_in_main(env):  # type: ignore
+            with futures.ThreadPoolExecutor(1) as executor:
+                create_pool_task = loop.run_in_executor(executor, start_pool_in_process)
+                process, process_status = await create_pool_task
 
         self._config.reset_pool_external_address(
             process_index, process_status.external_addresses[0]

{xoscar-0.6.1.dist-info → xoscar-0.6.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xoscar
-Version: 0.6.1
+Version: 0.6.2
 Summary: Python actor framework for heterogeneous computing.
 Home-page: http://github.com/xorbitsai/xoscar
 Author: Qin Xuye

{xoscar-0.6.1.dist-info → xoscar-0.6.2.dist-info}/RECORD

@@ -1,7 +1,7 @@
-xoscar-0.6.1.dist-info/RECORD,,
-xoscar-0.6.1.dist-info/WHEEL,sha256=e_QzmxBntUlJDFF5qAWxWPTItYuutkUby5BszMUIWnI,111
-xoscar-0.6.1.dist-info/top_level.txt,sha256=vYlqqY4Nys8Thm1hePIuUv8eQePdULVWMmt7lXtX_ZA,21
-xoscar-0.6.1.dist-info/METADATA,sha256=G1ypXuZl-ZjV48coc-vldRpkEeOq43k5BXIBowvxmo0,9076
+xoscar-0.6.2.dist-info/RECORD,,
+xoscar-0.6.2.dist-info/WHEEL,sha256=e_QzmxBntUlJDFF5qAWxWPTItYuutkUby5BszMUIWnI,111
+xoscar-0.6.2.dist-info/top_level.txt,sha256=vYlqqY4Nys8Thm1hePIuUv8eQePdULVWMmt7lXtX_ZA,21
+xoscar-0.6.2.dist-info/METADATA,sha256=W_mzFDfmy96PpO9ocfa9_wHgG2aa-geyIc_qJ4D-PHc,9076
 xoscar/_utils.pyx,sha256=6iqO4eTwEI-n9i39n_TKz7MWqbytMRnVNubdJ5egL6o,7279
 xoscar/backend.py,sha256=is436OPkZfSpQXaoqTRVta5eoye_pp45RFgCstAk2hU,1850
 xoscar/core.pxd,sha256=I_C2ka7XryyGnnAVXUVm8xfS1gtIrCs6X-9rswgOcUU,1317
@@ -10,7 +10,7 @@ xoscar/context.pxd,sha256=qKa0OyDPZtVymftSh447m-RzFZgmz8rGqQBa7qlauvc,725
 xoscar/batch.py,sha256=DpArS0L3WYJ_HVPG-6hSYEwoAFY1mY2-mlC4Jp5M_Dw,7872
 xoscar/nvutils.py,sha256=qmW4mKLU0WB2yCs198ccQOgLL02zB7Fsa-AotO3NOmg,20412
 xoscar/constants.py,sha256=QHHSREw6uWBBjQDCFqlNfTvBZgniJPGy42KSIsR8Fqw,787
-xoscar/_utils.cpython-310-darwin.so,sha256=G7A_8AL0avvF6Ez2nokyV3bDtBcuOAJW49QaCQgc9Xg,160688
+xoscar/_utils.cpython-310-darwin.so,sha256=gvVzg2It2b2yoL7ZbaPHzcq5kVRnyt0Qjh8c_am-2AI,160688
 xoscar/__init__.py,sha256=0zX8kKaio3ZIrlzB79WybcravMJw1OxPWjDspTgJFyQ,1608
 xoscar/api.py,sha256=3hztPoOxg8A_mlhWyWgVP7FMXG0PATA1TP4Rbaj7A-g,13327
 xoscar/utils.py,sha256=jUw6OICZUPBbmS1b3GE4vLctJf6fCKXrYtLtBuK-Oqc,16483
@@ -20,8 +20,8 @@ xoscar/context.pyx,sha256=8CdgPnWcE9eOp3N600WgDQ03MCi8P73eUOGcfV7Zksg,10942
 xoscar/errors.py,sha256=wBlQOKsXf0Fc4skN39tDie0YZT-VIAuLNRgoDl2pZcA,1241
 xoscar/core.pyx,sha256=phN-yYV0A0QI8WFi2jCu0nc4CnShTepfDi0V7ZrLYPY,22092
 xoscar/driver.py,sha256=498fowtJr6b3FE8FIOA_Tc1Vwx88nfZw7p0FxrML0h4,1372
-xoscar/context.cpython-310-darwin.so,sha256=iW3vrrt3TppY-wnq_lpXiTrUJHDnN9AmuLs2ac9blIg,200592
-xoscar/core.cpython-310-darwin.so,sha256=BB3Sy9cT8F8Y1eagVGQkh5Rlc67iff2jVTWMFzPt350,448704
+xoscar/context.cpython-310-darwin.so,sha256=YHhie4_b2dn8iZtxgGyeYwMGzM0K_SOxqrXwW3oktHY,200592
+xoscar/core.cpython-310-darwin.so,sha256=2uxow4CevQb3Az8U3WWZuBsu0QjeDtG3YRT7EANEaSc,448704
 xoscar/profiling.py,sha256=BC5OF0HzSaXv8V7w-y-B8r5gV5DgxHFoTEIF6jCMioQ,8015
 xoscar/_utils.pxd,sha256=5KYAL3jfPdejsHnrGGT2s--ZUX5SXznQWpHVSno429k,1157
 xoscar/metrics/__init__.py,sha256=9Badi7rxYikGm2dQiNCrj9GgMRBxwuR3JaEKcFZmfak,705
@@ -47,11 +47,11 @@ xoscar/serialization/cuda.py,sha256=iFUEnN4SiquBIhyieyOrfw3TnKnW-tU_vYgqOxO_DrA,
 xoscar/serialization/scipy.py,sha256=yOEi0NB8cqQ6e2UnCZ1w006RsB7T725tIL-DM_hNcsU,2482
 xoscar/serialization/aio.py,sha256=5DySPgDxU43ec7_5Ct44-Oqt7YNSJBfuf8VdQgQlChA,4731
 xoscar/serialization/core.pyx,sha256=bjR-zXGm9qersk7kYPzpjpMIxDl_Auur4BCubRfKmfA,29626
-xoscar/serialization/core.cpython-310-darwin.so,sha256=KZSkwjxb1OVAS9wmC936y2cHNRXaJ4YsvI2QD61KBIs,408208
+xoscar/serialization/core.cpython-310-darwin.so,sha256=Lx2pBNMODIaf27pDiuTAXdnRyiD2LcTlG7XyZ4X15Iw,408208
 xoscar/serialization/mlx.py,sha256=N_cvbTUBKc14XWYsPIMz4kDstyRN1DNhb4BVRgnQm8Y,1872
 xoscar/backends/config.py,sha256=4tZMiXAMMS8qQ4SX_LjONLtSQVfZTx3m-IK3EqbkYdk,5375
 xoscar/backends/allocate_strategy.py,sha256=tC1Nbq2tJohahUwd-zoRYHEDX65wyuX8tmeY45uWj_w,4845
-xoscar/backends/message.cpython-310-darwin.so,sha256=RJrc3_f95kNDF-7Nk_2lTe55oVNjwaGZ_d83EkBVj_M,387400
+xoscar/backends/message.cpython-310-darwin.so,sha256=tBsjEKzinkUPhCtP1SuN9qL3BBzpDMRVQa69RN6sEgo,387400
 xoscar/backends/__init__.py,sha256=VHEBQcUWM5bj027W8EUf9PiJUAP7JoMrRw3Tsvy5ySw,643
 xoscar/backends/core.py,sha256=EH-fHlV9x3bnruEHaUtGYO7osKLfLJ4AQHtuzA_mr2g,10857
 xoscar/backends/context.py,sha256=XfDPG2eDhAhE6hWBEkEsHTnyyOYN9R3houlMjAL7BFw,16329
@@ -61,7 +61,7 @@ xoscar/backends/pool.py,sha256=OezOhvvXAV3TpODhLHmJVgqCfowb3aA_fWZKPodm8bE,61003
 xoscar/backends/indigen/backend.py,sha256=znl_fZzWGEtLH8hZ9j9Kkf0fva25jEem2_KO7I1RVvc,1612
 xoscar/backends/indigen/__init__.py,sha256=tKHP5ClzedBRBpZsLRVErR3EUNbbDm4CY4u0rCFJr44,685
 xoscar/backends/indigen/driver.py,sha256=VGzkacYKykegW5qhCuhx01gdgBZEKJjNIyfNCnA6Nm8,952
-xoscar/backends/indigen/pool.py,sha256=v0Ps79W0WyeFSt2YxCJnh7q_1dnRQmo9887gcW3pNoc,17226
+xoscar/backends/indigen/pool.py,sha256=eeT29MDwou7UQrfY9B9Z7iY6TI7DlWvrZxg0euz96eA,18287
 xoscar/backends/test/backend.py,sha256=nv9WFhH5Bbq4Q1HB9yfpciZBaeHT4IQAtzugBWESrUY,1263
 xoscar/backends/test/__init__.py,sha256=j2ZfD6prD9WjUxRUDC7Eq5Z7N7TkL6fFr59oNyc_vY4,682
 xoscar/backends/test/pool.py,sha256=TW4X6J-92Pti66103poQBNDBznX6CBD3RLOc_zixjTo,7257