xoscar 0.4.5__tar.gz → 0.5.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of xoscar might be problematic. Click here for more details.

Files changed (85) hide show
  1. {xoscar-0.4.5 → xoscar-0.5.0}/PKG-INFO +32 -2
  2. {xoscar-0.4.5 → xoscar-0.5.0}/setup.py +1 -0
  3. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/_utils.pyx +6 -2
  4. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/communication/__init__.py +1 -1
  5. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/communication/socket.py +10 -2
  6. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/config.py +12 -0
  7. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/context.py +41 -8
  8. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/core.py +31 -5
  9. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/indigen/pool.py +16 -0
  10. xoscar-0.5.0/xoscar/backends/message.pyi +255 -0
  11. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/message.pyx +56 -9
  12. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/pool.py +30 -3
  13. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/router.py +86 -8
  14. xoscar-0.5.0/xoscar/collective/xoscar_pygloo.pyi +239 -0
  15. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/core.pxd +1 -0
  16. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/core.pyx +11 -5
  17. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/serialization/__init__.py +2 -2
  18. xoscar-0.5.0/xoscar/serialization/core.pyi +57 -0
  19. xoscar-0.5.0/xoscar/serialization/mlx.py +63 -0
  20. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar.egg-info/PKG-INFO +32 -2
  21. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar.egg-info/SOURCES.txt +4 -0
  22. {xoscar-0.4.5 → xoscar-0.5.0}/MANIFEST.in +0 -0
  23. {xoscar-0.4.5 → xoscar-0.5.0}/pyproject.toml +0 -0
  24. {xoscar-0.4.5 → xoscar-0.5.0}/setup.cfg +0 -0
  25. {xoscar-0.4.5 → xoscar-0.5.0}/versioneer.py +0 -0
  26. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/__init__.py +0 -0
  27. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/_utils.pxd +0 -0
  28. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/_version.py +0 -0
  29. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/aio/__init__.py +0 -0
  30. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/aio/base.py +0 -0
  31. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/aio/file.py +0 -0
  32. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/aio/lru.py +0 -0
  33. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/aio/parallelism.py +0 -0
  34. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/api.py +0 -0
  35. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backend.py +0 -0
  36. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/__init__.py +0 -0
  37. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/allocate_strategy.py +0 -0
  38. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/communication/base.py +0 -0
  39. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/communication/core.py +0 -0
  40. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/communication/dummy.py +0 -0
  41. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/communication/errors.py +0 -0
  42. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/communication/ucx.py +0 -0
  43. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/communication/utils.py +0 -0
  44. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/indigen/__init__.py +0 -0
  45. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/indigen/backend.py +0 -0
  46. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/indigen/driver.py +0 -0
  47. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/test/__init__.py +0 -0
  48. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/test/backend.py +0 -0
  49. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/backends/test/pool.py +0 -0
  50. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/batch.py +0 -0
  51. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/collective/__init__.py +0 -0
  52. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/collective/common.py +0 -0
  53. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/collective/core.py +0 -0
  54. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/collective/process_group.py +0 -0
  55. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/collective/utils.py +0 -0
  56. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/constants.py +0 -0
  57. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/context.pxd +0 -0
  58. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/context.pyx +0 -0
  59. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/debug.py +0 -0
  60. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/driver.py +0 -0
  61. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/errors.py +0 -0
  62. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/libcpp.pxd +0 -0
  63. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/metrics/__init__.py +0 -0
  64. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/metrics/api.py +0 -0
  65. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/metrics/backends/__init__.py +0 -0
  66. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/metrics/backends/console/__init__.py +0 -0
  67. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/metrics/backends/console/console_metric.py +0 -0
  68. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/metrics/backends/metric.py +0 -0
  69. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/metrics/backends/prometheus/__init__.py +0 -0
  70. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/metrics/backends/prometheus/prometheus_metric.py +0 -0
  71. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/nvutils.py +0 -0
  72. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/profiling.py +0 -0
  73. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/serialization/aio.py +0 -0
  74. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/serialization/core.pxd +0 -0
  75. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/serialization/core.pyx +0 -0
  76. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/serialization/cuda.py +0 -0
  77. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/serialization/exception.py +0 -0
  78. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/serialization/numpy.py +0 -0
  79. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/serialization/pyfury.py +0 -0
  80. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/serialization/scipy.py +0 -0
  81. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar/utils.py +0 -0
  82. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar.egg-info/dependency_links.txt +0 -0
  83. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar.egg-info/not-zip-safe +0 -0
  84. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar.egg-info/requires.txt +0 -0
  85. {xoscar-0.4.5 → xoscar-0.5.0}/xoscar.egg-info/top_level.txt +0 -0
@@ -1,6 +1,6 @@
1
- Metadata-Version: 2.1
1
+ Metadata-Version: 2.4
2
2
  Name: xoscar
3
- Version: 0.4.5
3
+ Version: 0.5.0
4
4
  Summary: Python actor framework for heterogeneous computing.
5
5
  Home-page: http://github.com/xorbitsai/xoscar
6
6
  Author: Qin Xuye
@@ -18,11 +18,41 @@ Classifier: Programming Language :: Python :: 3.12
18
18
  Classifier: Programming Language :: Python :: Implementation :: CPython
19
19
  Classifier: Topic :: Software Development :: Libraries
20
20
  Description-Content-Type: text/markdown
21
+ Requires-Dist: numpy>=1.14.0
22
+ Requires-Dist: pandas>=1.0.0
23
+ Requires-Dist: scipy>=1.0.0; sys_platform != "win32" or python_version >= "3.10"
24
+ Requires-Dist: scipy<=1.9.1,>=1.0.0; sys_platform == "win32" and python_version < "3.10"
25
+ Requires-Dist: cloudpickle>=1.5.0
26
+ Requires-Dist: psutil>=5.9.0
27
+ Requires-Dist: tblib>=1.7.0
28
+ Requires-Dist: uvloop>=0.14.0; sys_platform != "win32"
29
+ Requires-Dist: packaging
21
30
  Provides-Extra: dev
31
+ Requires-Dist: cython>=0.29; extra == "dev"
32
+ Requires-Dist: pytest>=3.5.0; extra == "dev"
33
+ Requires-Dist: pytest-cov>=2.5.0; extra == "dev"
34
+ Requires-Dist: pytest-timeout>=1.2.0; extra == "dev"
35
+ Requires-Dist: pytest-forked>=1.0; extra == "dev"
36
+ Requires-Dist: pytest-asyncio>=0.14.0; extra == "dev"
37
+ Requires-Dist: ipython>=6.5.0; extra == "dev"
38
+ Requires-Dist: sphinx; extra == "dev"
39
+ Requires-Dist: pydata-sphinx-theme>=0.3.0; extra == "dev"
40
+ Requires-Dist: sphinx-intl>=0.9.9; extra == "dev"
41
+ Requires-Dist: flake8>=3.8.0; extra == "dev"
42
+ Requires-Dist: black; extra == "dev"
22
43
  Provides-Extra: doc
44
+ Requires-Dist: ipython>=6.5.0; extra == "doc"
45
+ Requires-Dist: sphinx; extra == "doc"
46
+ Requires-Dist: pydata-sphinx-theme>=0.3.0; extra == "doc"
47
+ Requires-Dist: sphinx-intl>=0.9.9; extra == "doc"
23
48
  Provides-Extra: extra
49
+ Requires-Dist: pyarrow>=5.0.0; extra == "extra"
24
50
  Provides-Extra: kubernetes
51
+ Requires-Dist: kubernetes>=10.0.0; extra == "kubernetes"
25
52
  Provides-Extra: ray
53
+ Requires-Dist: xoscar_ray>=0.0.1; extra == "ray"
54
+ Dynamic: description
55
+ Dynamic: description-content-type
26
56
 
27
57
  <div align="center">
28
58
  <img width="77%" alt="" src="https://raw.githubusercontent.com/xprobe-inc/xoscar/main/doc/source/_static/Xoscar.svg"><br>
@@ -290,6 +290,7 @@ class CMakeBuild(build_ext):
290
290
  f"-DBUILD_TMP_DIR={build_temp}",
291
291
  f"-DPYTHON_PATH={sys.executable}",
292
292
  f"-DCMAKE_BUILD_TYPE={cfg}", # not used on MSVC, but no harm
293
+ "-DCMAKE_POLICY_VERSION_MINIMUM=3.10",
293
294
  ]
294
295
  build_args = []
295
296
  # Adding CMake arguments set as environment variable
@@ -208,28 +208,32 @@ def create_actor_ref(*args, **kwargs):
208
208
 
209
209
  address = to_str(kwargs.pop('address', None))
210
210
  uid = kwargs.pop('uid', None)
211
+ proxy_addresses = kwargs.pop("proxy_addresses", None)
211
212
 
212
213
  if kwargs:
213
214
  raise ValueError('Only `address` or `uid` keywords are supported')
214
215
 
215
- if len(args) == 2:
216
+ if 2 <= len(args) <= 3:
216
217
  if address:
217
218
  raise ValueError('address has been specified')
218
219
  address = to_str(args[0])
219
220
  uid = args[1]
221
+ if len(args) == 3:
222
+ proxy_addresses = args[2]
220
223
  elif len(args) == 1:
221
224
  tp0 = type(args[0])
222
225
  if tp0 is ActorRef or tp0 is LocalActorRef:
223
226
  existing_ref = <ActorRef>(args[0])
224
227
  uid = existing_ref.uid
225
228
  address = to_str(address or existing_ref.address)
229
+ proxy_addresses = existing_ref.proxy_addresses
226
230
  else:
227
231
  uid = args[0]
228
232
 
229
233
  if uid is None:
230
234
  raise ValueError('Actor uid should be provided')
231
235
 
232
- return ActorRef(address, uid)
236
+ return ActorRef(address, uid, proxy_addresses=proxy_addresses)
233
237
 
234
238
 
235
239
  cdef class Timer:
@@ -13,7 +13,7 @@
13
13
  # See the License for the specific language governing permissions and
14
14
  # limitations under the License.
15
15
 
16
- from .base import Channel, Client, Server
16
+ from .base import Channel, ChannelType, Client, Server
17
17
  from .core import gen_local_address, get_client_type, get_server_type
18
18
  from .dummy import DummyChannel, DummyClient, DummyServer
19
19
  from .socket import (
@@ -315,7 +315,11 @@ class SocketClient(Client):
315
315
  except asyncio.TimeoutError:
316
316
  raise ConnectionError("connect timeout")
317
317
  channel = SocketChannel(
318
- reader, writer, local_address=local_address, dest_address=dest_address
318
+ reader,
319
+ writer,
320
+ local_address=local_address,
321
+ dest_address=dest_address,
322
+ channel_type=ChannelType.remote,
319
323
  )
320
324
  return SocketClient(local_address, dest_address, channel)
321
325
 
@@ -431,6 +435,10 @@ class UnixSocketClient(Client):
431
435
  "Cannot connect unix socket due to file not exists"
432
436
  )
433
437
  channel = SocketChannel(
434
- reader, writer, local_address=local_address, dest_address=dest_address
438
+ reader,
439
+ writer,
440
+ local_address=local_address,
441
+ dest_address=dest_address,
442
+ channel_type=ChannelType.ipc,
435
443
  )
436
444
  return UnixSocketClient(local_address, dest_address, channel)
@@ -33,6 +33,8 @@ class ActorPoolConfig:
33
33
  self._conf["metrics"] = dict()
34
34
  if "comm" not in self._conf:
35
35
  self._conf["comm"] = dict()
36
+ if "proxy" not in self._conf:
37
+ self._conf["proxy"] = dict()
36
38
 
37
39
  @property
38
40
  def n_pool(self):
@@ -143,3 +145,13 @@ class ActorPoolConfig:
143
145
 
144
146
  def get_comm_config(self) -> dict:
145
147
  return self._conf["comm"]
148
+
149
+ def get_proxy_config(self) -> dict:
150
+ return self._conf["proxy"]
151
+
152
+ def add_proxy_config(self, proxy_config: dict[str, str] | None):
153
+ if proxy_config:
154
+ self._conf["proxy"].update(proxy_config)
155
+
156
+ def remove_proxy(self, from_addr: str):
157
+ del self._conf["proxy"][from_addr]
@@ -72,15 +72,24 @@ class IndigenActorContext(BaseActorContext):
72
72
  self._caller.cancel_tasks()
73
73
 
74
74
  async def _call(
75
- self, address: str, message: _MessageBase, wait: bool = True
75
+ self,
76
+ address: str,
77
+ message: _MessageBase,
78
+ wait: bool = True,
79
+ proxy_addresses: list[str] | None = None,
76
80
  ) -> Union[ResultMessage, ErrorMessage, asyncio.Future]:
77
81
  return await self._caller.call(
78
- Router.get_instance_or_empty(), address, message, wait=wait
82
+ Router.get_instance_or_empty(),
83
+ address,
84
+ message,
85
+ wait=wait,
86
+ proxy_addresses=proxy_addresses,
79
87
  )
80
88
 
81
89
  async def _call_with_client(
82
90
  self, client: Client, message: _MessageBase, wait: bool = True
83
91
  ) -> Union[ResultMessage, ErrorMessage, asyncio.Future]:
92
+ # NOTE: used by copyto, cannot support proxy
84
93
  return await self._caller.call_with_client(client, message, wait)
85
94
 
86
95
  async def _call_send_buffers(
@@ -146,7 +155,12 @@ class IndigenActorContext(BaseActorContext):
146
155
  message = HasActorMessage(
147
156
  new_message_id(), actor_ref, protocol=DEFAULT_PROTOCOL
148
157
  )
149
- future = await self._call(actor_ref.address, message, wait=False)
158
+ future = await self._call(
159
+ actor_ref.address,
160
+ message,
161
+ wait=False,
162
+ proxy_addresses=actor_ref.proxy_addresses,
163
+ )
150
164
  result = await self._wait(future, actor_ref.address, message) # type: ignore
151
165
  return self._process_result_message(result)
152
166
 
@@ -154,7 +168,12 @@ class IndigenActorContext(BaseActorContext):
154
168
  message = DestroyActorMessage(
155
169
  new_message_id(), actor_ref, protocol=DEFAULT_PROTOCOL
156
170
  )
157
- future = await self._call(actor_ref.address, message, wait=False)
171
+ future = await self._call(
172
+ actor_ref.address,
173
+ message,
174
+ wait=False,
175
+ proxy_addresses=actor_ref.proxy_addresses,
176
+ )
158
177
  result = await self._wait(future, actor_ref.address, message) # type: ignore
159
178
  return self._process_result_message(result)
160
179
 
@@ -168,7 +187,7 @@ class IndigenActorContext(BaseActorContext):
168
187
  protocol=DEFAULT_PROTOCOL,
169
188
  )
170
189
  main_address = self._process_result_message(
171
- await self._call(actor_ref.address, control_message) # type: ignore
190
+ await self._call(actor_ref.address, control_message, proxy_addresses=actor_ref.proxy_addresses) # type: ignore
172
191
  )
173
192
  real_actor_ref = await self.actor_ref(actor_ref)
174
193
  if real_actor_ref.address == main_address:
@@ -182,7 +201,9 @@ class IndigenActorContext(BaseActorContext):
182
201
  protocol=DEFAULT_PROTOCOL,
183
202
  )
184
203
  # stop server
185
- result = await self._call(main_address, stop_message)
204
+ result = await self._call(
205
+ main_address, stop_message, proxy_addresses=actor_ref.proxy_addresses
206
+ )
186
207
  return self._process_result_message(result) # type: ignore
187
208
 
188
209
  async def actor_ref(self, *args, **kwargs):
@@ -194,7 +215,12 @@ class IndigenActorContext(BaseActorContext):
194
215
  message = ActorRefMessage(
195
216
  new_message_id(), actor_ref, protocol=DEFAULT_PROTOCOL
196
217
  )
197
- future = await self._call(actor_ref.address, message, wait=False)
218
+ future = await self._call(
219
+ actor_ref.address,
220
+ message,
221
+ wait=False,
222
+ proxy_addresses=actor_ref.proxy_addresses,
223
+ )
198
224
  result = await self._wait(future, actor_ref.address, message)
199
225
  res = self._process_result_message(result)
200
226
  if res.address != connect_addr:
@@ -225,7 +251,12 @@ class IndigenActorContext(BaseActorContext):
225
251
  actor_ref.address,
226
252
  ):
227
253
  detect_cycle_send(send_message, wait_response)
228
- future = await self._call(actor_ref.address, send_message, wait=False)
254
+ future = await self._call(
255
+ actor_ref.address,
256
+ send_message,
257
+ wait=False,
258
+ proxy_addresses=actor_ref.proxy_addresses,
259
+ )
229
260
  if wait_response:
230
261
  result = await self._wait(future, actor_ref.address, send_message) # type: ignore
231
262
  return self._process_result_message(result)
@@ -315,6 +346,8 @@ class IndigenActorContext(BaseActorContext):
315
346
  async def _get_client(self, address: str) -> Client:
316
347
  router = Router.get_instance()
317
348
  assert router is not None, "`copy_to` can only be used inside pools"
349
+ if router.get_proxy(address):
350
+ raise RuntimeError("Cannot run `copy_to` when enabling proxy")
318
351
  return await self._get_copy_to_client(router, address)
319
352
 
320
353
  async def copy_to_buffers(
@@ -26,8 +26,15 @@ from typing import Type, Union
26
26
  from .._utils import Timer
27
27
  from ..errors import ServerClosed
28
28
  from ..profiling import get_profiling_data
29
- from .communication import Client, UCXClient
30
- from .message import DeserializeMessageFailed, ErrorMessage, ResultMessage, _MessageBase
29
+ from .communication import ChannelType, Client, UCXClient
30
+ from .message import (
31
+ DeserializeMessageFailed,
32
+ ErrorMessage,
33
+ ForwardMessage,
34
+ MessageType,
35
+ ResultMessage,
36
+ _MessageBase,
37
+ )
31
38
  from .router import Router
32
39
 
33
40
  ResultMessageType = Union[ResultMessage, ErrorMessage]
@@ -67,8 +74,15 @@ class ActorCallerThreadLocal:
67
74
  self._listen_client(client)
68
75
  return client
69
76
 
70
- async def get_client(self, router: Router, dest_address: str) -> Client:
71
- client = await router.get_client(dest_address, from_who=self)
77
+ async def get_client(
78
+ self,
79
+ router: Router,
80
+ dest_address: str,
81
+ proxy_addresses: list[str] | None = None,
82
+ ) -> Client:
83
+ client = await router.get_client(
84
+ dest_address, from_who=self, proxy_addresses=proxy_addresses
85
+ )
72
86
  self._listen_client(client)
73
87
  return client
74
88
 
@@ -191,8 +205,20 @@ class ActorCallerThreadLocal:
191
205
  dest_address: str,
192
206
  message: _MessageBase,
193
207
  wait: bool = True,
208
+ proxy_addresses: list[str] | None = None,
194
209
  ) -> ResultMessage | ErrorMessage | asyncio.Future:
195
- client = await self.get_client(router, dest_address)
210
+ client = await self.get_client(
211
+ router, dest_address, proxy_addresses=proxy_addresses
212
+ )
213
+ if (
214
+ client.channel_type == ChannelType.remote
215
+ and client.dest_address != dest_address
216
+ and message.message_type != MessageType.control
217
+ ):
218
+ # wrap message with forward message
219
+ message = ForwardMessage(
220
+ message_id=message.message_id, address=dest_address, raw_message=message
221
+ )
196
222
  return await self.call_with_client(client, message, wait)
197
223
 
198
224
  async def stop(self):
@@ -16,6 +16,7 @@
16
16
  from __future__ import annotations
17
17
 
18
18
  import asyncio
19
+ import atexit
19
20
  import concurrent.futures as futures
20
21
  import configparser
21
22
  import contextlib
@@ -29,6 +30,7 @@ import sys
29
30
  import threading
30
31
  import uuid
31
32
  from dataclasses import dataclass
33
+ from multiprocessing import util
32
34
  from types import TracebackType
33
35
  from typing import List, Optional
34
36
 
@@ -79,6 +81,19 @@ logger = logging.getLogger(__name__)
79
81
  _init_main_suspended_local = threading.local()
80
82
 
81
83
 
84
+ def _terminate_children():
85
+ for c in multiprocessing.active_children():
86
+ try:
87
+ c.terminate()
88
+ except Exception:
89
+ pass
90
+
91
+
92
+ if util:
93
+ # Import multiprocessing.util to register _exit_function at exit.
94
+ atexit.register(_terminate_children)
95
+
96
+
82
97
  def _patch_spawn_get_preparation_data():
83
98
  try:
84
99
  from multiprocessing import spawn as mp_spawn
@@ -309,6 +324,7 @@ class MainActorPool(MainActorPoolBase):
309
324
  raise
310
325
  finally:
311
326
  status_queue.put(process_status)
327
+ status_queue.cancel_join_thread()
312
328
  await pool.join()
313
329
 
314
330
  async def append_sub_pool(
@@ -0,0 +1,255 @@
1
+ # Copyright 2022-2023 XProbe Inc.
2
+ # derived from copyright 1999-2022 Alibaba Group Holding Ltd.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ from __future__ import annotations
17
+
18
+ from enum import Enum
19
+ from types import TracebackType
20
+ from typing import Any, List, Type
21
+
22
+ from ..core import ActorRef, BufferRef
23
+
24
+ DEFAULT_PROTOCOL: int = 0
25
+
26
+ class MessageType(Enum):
27
+ control = 0
28
+ result = 1
29
+ error = 2
30
+ create_actor = 3
31
+ destroy_actor = 4
32
+ has_actor = 5
33
+ actor_ref = 6
34
+ send = 7
35
+ tell = 8
36
+ cancel = 9
37
+ copy_to_buffers = 10
38
+ copy_to_fileobjs = 11
39
+ forward = 12 # message forwarded to other pool
40
+
41
+ class ControlMessageType(Enum):
42
+ stop = 0
43
+ restart = 1
44
+ sync_config = 2
45
+ get_config = 3
46
+ wait_pool_recovered = 4
47
+ add_sub_pool_actor = 5
48
+ # indicate that the following data will be used for copy_to
49
+ switch_to_copy_to = 6
50
+
51
+ class _MessageBase:
52
+ message_type: MessageType
53
+ protocol: int
54
+ message_id: bytes
55
+ message_trace: list
56
+ profiling_context: Any
57
+
58
+ def __init__(
59
+ self,
60
+ message_id: bytes | None = None,
61
+ protocol: int = DEFAULT_PROTOCOL,
62
+ message_trace: list | None = None,
63
+ profiling_context: Any = None,
64
+ ): ...
65
+ def __repr__(self): ...
66
+
67
+ class CopyToBuffersMessage(_MessageBase):
68
+ message_type = MessageType.copy_to_buffers
69
+
70
+ content: object
71
+
72
+ def __int__(
73
+ self,
74
+ message_id: bytes | None = None,
75
+ content: object = None,
76
+ protocol: int = DEFAULT_PROTOCOL,
77
+ message_trace: list | None = None,
78
+ ): ...
79
+
80
+ class CopyToFileObjectsMessage(CopyToBuffersMessage):
81
+ message_type = MessageType.copy_to_fileobjs
82
+
83
+ class ControlMessage(_MessageBase):
84
+ message_type = MessageType.control
85
+
86
+ address: str
87
+ control_message_type: ControlMessageType
88
+ content: Any
89
+
90
+ def __init__(
91
+ self,
92
+ message_id: bytes | None = None,
93
+ address: str | None = None,
94
+ control_message_type: ControlMessageType | None = None,
95
+ content: Any = None,
96
+ protocol: int = DEFAULT_PROTOCOL,
97
+ message_trace: list | None = None,
98
+ ): ...
99
+
100
+ class ResultMessage(_MessageBase):
101
+ message_type = MessageType.result
102
+
103
+ result: Any
104
+
105
+ def __init__(
106
+ self,
107
+ message_id: bytes | None = None,
108
+ result: Any = None,
109
+ protocol: int = DEFAULT_PROTOCOL,
110
+ message_trace: list | None = None,
111
+ profiling_context: Any = None,
112
+ ): ...
113
+
114
+ class ErrorMessage(_MessageBase):
115
+ message_type = MessageType.error
116
+
117
+ address: str
118
+ pid: int
119
+ error_type: Type
120
+ error: BaseException
121
+ traceback: TracebackType
122
+
123
+ def __init__(
124
+ self,
125
+ message_id: bytes | None = None,
126
+ address: str | None = None,
127
+ pid: int = -1,
128
+ error_type: Type[BaseException] | None = None,
129
+ error: BaseException | None = None,
130
+ traceback: TracebackType | None = None,
131
+ protocol: int = DEFAULT_PROTOCOL,
132
+ message_trace: list | None = None,
133
+ ): ...
134
+ def as_instanceof_cause(self) -> BaseException: ...
135
+
136
+ class CreateActorMessage(_MessageBase):
137
+ message_type = MessageType.create_actor
138
+
139
+ actor_cls: Type
140
+ actor_id: bytes
141
+ args: tuple
142
+ kwargs: dict
143
+ allocate_strategy: Any
144
+ from_main: bool
145
+
146
+ def __init__(
147
+ self,
148
+ message_id: bytes | None = None,
149
+ actor_cls: Type | None = None,
150
+ actor_id: bytes | None = None,
151
+ args: tuple | None = None,
152
+ kwargs: dict | None = None,
153
+ allocate_strategy: Any = None,
154
+ from_main: bool = False,
155
+ protocol: int = DEFAULT_PROTOCOL,
156
+ message_trace: list | None = None,
157
+ ): ...
158
+
159
+ class DestroyActorMessage(_MessageBase):
160
+ message_type = MessageType.destroy_actor
161
+
162
+ actor_ref: ActorRef
163
+ from_main: bool
164
+
165
+ def __init__(
166
+ self,
167
+ message_id: bytes | None = None,
168
+ actor_ref: ActorRef = None,
169
+ from_main: bool = False,
170
+ protocol: int = DEFAULT_PROTOCOL,
171
+ message_trace: list | None = None,
172
+ ): ...
173
+
174
+ class HasActorMessage(_MessageBase):
175
+ message_type = MessageType.has_actor
176
+
177
+ actor_ref: ActorRef
178
+
179
+ def __init__(
180
+ self,
181
+ message_id: bytes | None = None,
182
+ actor_ref: ActorRef = None,
183
+ protocol: int = DEFAULT_PROTOCOL,
184
+ message_trace: list | None = None,
185
+ ): ...
186
+
187
+ class ActorRefMessage(_MessageBase):
188
+ message_type = MessageType.actor_ref
189
+
190
+ actor_ref: ActorRef
191
+
192
+ def __init__(
193
+ self,
194
+ message_id: bytes | None = None,
195
+ actor_ref: ActorRef = None,
196
+ protocol: int = DEFAULT_PROTOCOL,
197
+ message_trace: list | None = None,
198
+ ): ...
199
+
200
+ class SendMessage(_MessageBase):
201
+ message_type = MessageType.send
202
+
203
+ actor_ref: ActorRef
204
+ content: Any
205
+
206
+ def __init__(
207
+ self,
208
+ message_id: bytes | None = None,
209
+ actor_ref: ActorRef = None,
210
+ content: object = None,
211
+ protocol: int = DEFAULT_PROTOCOL,
212
+ message_trace: list | None = None,
213
+ profiling_context: Any = None,
214
+ ): ...
215
+
216
+ class TellMessage(SendMessage):
217
+ message_type = MessageType.tell
218
+
219
+ class CancelMessage(_MessageBase):
220
+ message_type = MessageType.cancel
221
+
222
+ address: str
223
+ cancel_message_id: bytes
224
+
225
+ def __init__(
226
+ self,
227
+ message_id: bytes | None = None,
228
+ address: str | None = None,
229
+ cancel_message_id: bytes | None = None,
230
+ protocol: int = DEFAULT_PROTOCOL,
231
+ message_trace: list | None = None,
232
+ ): ...
233
+
234
+ class ForwardMessage(_MessageBase):
235
+ message_type = MessageType.forward
236
+
237
+ address: str
238
+ raw_message: _MessageBase
239
+
240
+ def __init__(
241
+ self,
242
+ message_id: bytes | None = None,
243
+ address: str | None = None,
244
+ raw_message: _MessageBase | None = None,
245
+ protocol: int = DEFAULT_PROTOCOL,
246
+ message_trace: list | None = None,
247
+ ): ...
248
+
249
+ class DeserializeMessageFailed(RuntimeError):
250
+ message_id: bytes
251
+
252
+ def __init__(self, message_id: bytes): ...
253
+ def __str__(self): ...
254
+
255
+ def new_message_id() -> bytes: ...