xinference 1.4.0__py3-none-any.whl → 1.4.1__py3-none-any.whl

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 xinference might be problematic. Click here for more details.

Files changed (59) hide show
  1. xinference/_compat.py +1 -0
  2. xinference/_version.py +3 -3
  3. xinference/api/restful_api.py +4 -0
  4. xinference/core/model.py +23 -3
  5. xinference/core/supervisor.py +6 -0
  6. xinference/core/worker.py +54 -11
  7. xinference/model/llm/__init__.py +4 -2
  8. xinference/model/llm/core.py +1 -0
  9. xinference/model/llm/llama_cpp/core.py +6 -1
  10. xinference/model/llm/llm_family.json +117 -1
  11. xinference/model/llm/llm_family_modelscope.json +125 -1
  12. xinference/model/llm/reasoning_parser.py +3 -3
  13. xinference/model/llm/sglang/core.py +111 -13
  14. xinference/model/llm/transformers/core.py +1 -0
  15. xinference/model/llm/transformers/deepseek_vl.py +1 -1
  16. xinference/model/llm/transformers/deepseek_vl2.py +287 -0
  17. xinference/model/llm/utils.py +26 -14
  18. xinference/model/llm/vllm/core.py +149 -8
  19. xinference/model/llm/vllm/distributed_executor.py +314 -0
  20. xinference/model/rerank/core.py +16 -11
  21. xinference/thirdparty/deepseek_vl2/__init__.py +31 -0
  22. xinference/thirdparty/deepseek_vl2/models/__init__.py +26 -0
  23. xinference/thirdparty/deepseek_vl2/models/configuration_deepseek.py +210 -0
  24. xinference/thirdparty/deepseek_vl2/models/conversation.py +310 -0
  25. xinference/thirdparty/deepseek_vl2/models/modeling_deepseek.py +1975 -0
  26. xinference/thirdparty/deepseek_vl2/models/modeling_deepseek_vl_v2.py +697 -0
  27. xinference/thirdparty/deepseek_vl2/models/processing_deepseek_vl_v2.py +675 -0
  28. xinference/thirdparty/deepseek_vl2/models/siglip_vit.py +661 -0
  29. xinference/thirdparty/deepseek_vl2/serve/__init__.py +0 -0
  30. xinference/thirdparty/deepseek_vl2/serve/app_modules/__init__.py +0 -0
  31. xinference/thirdparty/deepseek_vl2/serve/app_modules/gradio_utils.py +83 -0
  32. xinference/thirdparty/deepseek_vl2/serve/app_modules/overwrites.py +81 -0
  33. xinference/thirdparty/deepseek_vl2/serve/app_modules/presets.py +115 -0
  34. xinference/thirdparty/deepseek_vl2/serve/app_modules/utils.py +333 -0
  35. xinference/thirdparty/deepseek_vl2/serve/assets/Kelpy-Codos.js +100 -0
  36. xinference/thirdparty/deepseek_vl2/serve/assets/avatar.png +0 -0
  37. xinference/thirdparty/deepseek_vl2/serve/assets/custom.css +355 -0
  38. xinference/thirdparty/deepseek_vl2/serve/assets/custom.js +22 -0
  39. xinference/thirdparty/deepseek_vl2/serve/assets/favicon.ico +0 -0
  40. xinference/thirdparty/deepseek_vl2/serve/assets/simsun.ttc +0 -0
  41. xinference/thirdparty/deepseek_vl2/serve/inference.py +197 -0
  42. xinference/thirdparty/deepseek_vl2/utils/__init__.py +18 -0
  43. xinference/thirdparty/deepseek_vl2/utils/io.py +80 -0
  44. xinference/web/ui/build/asset-manifest.json +3 -3
  45. xinference/web/ui/build/index.html +1 -1
  46. xinference/web/ui/build/static/js/{main.3cea968e.js → main.5ca4eea1.js} +3 -3
  47. xinference/web/ui/build/static/js/main.5ca4eea1.js.map +1 -0
  48. xinference/web/ui/node_modules/.cache/babel-loader/0f0967acaec5df1d45b80010949c258d64297ebbb0f44b8bb3afcbd45c6f0ec4.json +1 -0
  49. xinference/web/ui/node_modules/.cache/babel-loader/68249645124f37d01eef83b1d897e751f895bea919b6fb466f907c1f87cebc84.json +1 -0
  50. {xinference-1.4.0.dist-info → xinference-1.4.1.dist-info}/METADATA +4 -4
  51. {xinference-1.4.0.dist-info → xinference-1.4.1.dist-info}/RECORD +56 -31
  52. xinference/web/ui/build/static/js/main.3cea968e.js.map +0 -1
  53. xinference/web/ui/node_modules/.cache/babel-loader/7f59e45e3f268ab8a4788b6fb024cf8dab088736dff22f5a3a39c122a83ab930.json +0 -1
  54. xinference/web/ui/node_modules/.cache/babel-loader/dcd60488509450bfff37bfff56de2c096d51de17dd00ec60d4db49c8b483ada1.json +0 -1
  55. /xinference/web/ui/build/static/js/{main.3cea968e.js.LICENSE.txt → main.5ca4eea1.js.LICENSE.txt} +0 -0
  56. {xinference-1.4.0.dist-info → xinference-1.4.1.dist-info}/LICENSE +0 -0
  57. {xinference-1.4.0.dist-info → xinference-1.4.1.dist-info}/WHEEL +0 -0
  58. {xinference-1.4.0.dist-info → xinference-1.4.1.dist-info}/entry_points.txt +0 -0
  59. {xinference-1.4.0.dist-info → xinference-1.4.1.dist-info}/top_level.txt +0 -0
@@ -13,12 +13,16 @@
13
13
  # limitations under the License.
14
14
 
15
15
  import asyncio
16
+ import itertools
16
17
  import json
17
18
  import logging
18
19
  import multiprocessing
19
20
  import os
21
+ import sys
22
+ import threading
20
23
  import time
21
24
  import uuid
25
+ from functools import partial
22
26
  from typing import (
23
27
  TYPE_CHECKING,
24
28
  Any,
@@ -27,10 +31,13 @@ from typing import (
27
31
  List,
28
32
  Optional,
29
33
  Tuple,
34
+ Type,
30
35
  TypedDict,
31
36
  Union,
32
37
  )
33
38
 
39
+ import xoscar as xo
40
+
34
41
  from ....types import (
35
42
  ChatCompletion,
36
43
  ChatCompletionChunk,
@@ -73,6 +80,7 @@ class VLLMModelConfig(TypedDict, total=False):
73
80
  guided_decoding_backend: Optional[str]
74
81
  scheduling_policy: Optional[str]
75
82
  reasoning_content: bool
83
+ model_quantization: Optional[str]
76
84
 
77
85
 
78
86
  class VLLMGenerateConfig(TypedDict, total=False):
@@ -161,6 +169,7 @@ if VLLM_INSTALLED and vllm.__version__ >= "0.3.0":
161
169
  VLLM_SUPPORTED_CHAT_MODELS.append("QwQ-32B")
162
170
  VLLM_SUPPORTED_CHAT_MODELS.append("marco-o1")
163
171
  VLLM_SUPPORTED_CHAT_MODELS.append("deepseek-r1-distill-qwen")
172
+ VLLM_SUPPORTED_CHAT_MODELS.append("fin-r1")
164
173
 
165
174
  if VLLM_INSTALLED and vllm.__version__ >= "0.3.2":
166
175
  VLLM_SUPPORTED_CHAT_MODELS.append("gemma-it")
@@ -248,15 +257,59 @@ class VLLMModel(LLM):
248
257
  self.lora_modules = peft_model
249
258
  self.lora_requests: List[LoRARequest] = []
250
259
  self._xavier_config = None
260
+ # distributed inference
261
+ self._device_count = None
262
+ self._address = model_config.pop("address", None) # type: ignore
263
+ self._n_worker = model_config.pop("n_worker", 1) # type: ignore
264
+ self._shard = model_config.pop("shard", 0) # type: ignore
265
+ self._driver_info = model_config.pop("driver_info", None) # type: ignore
266
+ self._loading_thread: Optional[threading.Thread] = None
267
+ self._loading_error = None
268
+ # variables used for distributed inference and multiple GPUs
269
+ self._pool_addresses = None
270
+ self._worker_addresses: Optional[Dict[int, List[str]]] = None
271
+ self._all_worker_ready: Optional[threading.Event] = None
272
+ # used to call async
273
+ self._loop = None
251
274
 
252
275
  def set_xavier_config(self, value: Optional[Dict]):
253
276
  self._xavier_config = value # type: ignore
254
277
 
278
+ def set_worker_addresses(self, shard: int, worker_addresses: List[str]):
279
+ assert self._worker_addresses is not None
280
+ self._worker_addresses[shard] = worker_addresses
281
+ if (
282
+ self._all_worker_ready is not None
283
+ and len(self._worker_addresses) == self._n_worker
284
+ ):
285
+ self._all_worker_ready.set()
286
+
287
+ @property
288
+ def driver_info(self) -> Optional[dict]:
289
+ return self._driver_info
290
+
291
+ @property
292
+ def need_create_pools(self):
293
+ return True
294
+
295
+ def set_pool_addresses(self, pool_addresses: List[str]):
296
+ self._pool_addresses = pool_addresses # type: ignore
297
+
298
+ def get_pool_addresses(self) -> Optional[List[str]]:
299
+ return self._pool_addresses
300
+
301
+ def set_loop(self, loop: asyncio.AbstractEventLoop):
302
+ # loop will be passed into XinferenceDistributedExecutor,
303
+ # to call aynsc method with asyncio.run_coroutine_threadsafe
304
+ self._loop = loop # type: ignore
305
+
255
306
  def load(self):
256
307
  try:
257
308
  import vllm
309
+ from vllm.config import VllmConfig
258
310
  from vllm.engine.arg_utils import AsyncEngineArgs
259
311
  from vllm.engine.async_llm_engine import AsyncLLMEngine
312
+ from vllm.executor.executor_base import ExecutorBase
260
313
  from vllm.lora.request import LoRARequest
261
314
  except ImportError:
262
315
  error_message = "Failed to import module 'vllm'"
@@ -275,6 +328,7 @@ class VLLMModel(LLM):
275
328
  # we need to set it to fork to make cupy NCCL work
276
329
  multiprocessing.set_start_method("fork", force=True)
277
330
 
331
+ self._device_count = self._get_cuda_count()
278
332
  self._model_config = self._sanitize_model_config(self._model_config)
279
333
  reasoning_content = self._model_config.pop("reasoning_content")
280
334
 
@@ -320,6 +374,83 @@ class VLLMModel(LLM):
320
374
  self._engine = XavierEngine.from_engine_args(
321
375
  engine_args, xavier_config=self._xavier_config
322
376
  )
377
+ elif self._n_worker > 1 or (
378
+ self._device_count > 1 and vllm.__version__ >= "0.7.0"
379
+ ):
380
+ from .distributed_executor import XinferenceDistributedExecutor
381
+
382
+ # model across multiple workers or GPUs
383
+ engine_args = AsyncEngineArgs(
384
+ model=self.model_path,
385
+ enable_lora=enable_lora,
386
+ max_loras=max_loras,
387
+ **self._model_config,
388
+ )
389
+
390
+ assert self._loop is not None
391
+ self._worker_addresses = {}
392
+
393
+ def _load():
394
+ try:
395
+ assert self._pool_addresses
396
+
397
+ if self._shard > 0:
398
+ assert self._driver_info
399
+ address = self._driver_info["address"]
400
+
401
+ coro = xo.actor_ref(address, self.raw_model_uid)
402
+ model_ref = asyncio.run_coroutine_threadsafe(
403
+ coro, self._loop
404
+ ).result()
405
+ coro = model_ref.set_worker_addresses(
406
+ self._shard, self._pool_addresses
407
+ )
408
+ asyncio.run_coroutine_threadsafe(coro, self._loop).result()
409
+ else:
410
+ self.set_worker_addresses(0, self._pool_addresses)
411
+ self._driver_info = {"address": self._address}
412
+
413
+ if self._n_worker > 1:
414
+ self._all_worker_ready = threading.Event()
415
+ # if model across workers, wait for other workers ready
416
+ self._all_worker_ready.wait()
417
+
418
+ # gather all worker addresses
419
+ worker_addresses = list(
420
+ itertools.chain(
421
+ *[
422
+ self._worker_addresses[shard]
423
+ for shard in range(self._n_worker)
424
+ ]
425
+ )
426
+ )
427
+ assert worker_addresses
428
+ loop = self._loop
429
+
430
+ class XinferenceAsyncLLMEngine(AsyncLLMEngine):
431
+ @classmethod
432
+ def _get_executor_cls(
433
+ cls, engine_config: VllmConfig
434
+ ) -> Type[ExecutorBase]:
435
+ return partial( # type: ignore
436
+ XinferenceDistributedExecutor,
437
+ pool_addresses=worker_addresses,
438
+ n_worker=self._n_worker,
439
+ loop=loop,
440
+ )
441
+
442
+ self._engine = XinferenceAsyncLLMEngine.from_engine_args(
443
+ engine_args
444
+ )
445
+ except:
446
+ logger.exception("Creating vllm engine failed")
447
+ self._loading_error = sys.exc_info()
448
+
449
+ self._loading_thread = threading.Thread(target=_load)
450
+ self._loading_thread.start()
451
+ # wait some time for init finish
452
+ if self._shard == 0:
453
+ self._loading_thread.join(1)
323
454
  else:
324
455
  engine_args = AsyncEngineArgs(
325
456
  model=self.model_path,
@@ -332,7 +463,14 @@ class VLLMModel(LLM):
332
463
  self._check_health_task = None
333
464
  if hasattr(self._engine, "check_health"):
334
465
  # vLLM introduced `check_health` since v0.4.1
335
- self._check_health_task = asyncio.create_task(self._check_healthy())
466
+ self._check_health_task = self._loop.create_task(self._check_healthy())
467
+
468
+ def wait_for_load(self):
469
+ if self._loading_thread:
470
+ self._loading_thread.join()
471
+ if self._loading_error:
472
+ _, err, tb = self._loading_error
473
+ raise err.with_traceback(tb)
336
474
 
337
475
  def stop(self):
338
476
  # though the vLLM engine will shutdown when deleted,
@@ -341,9 +479,10 @@ class VLLMModel(LLM):
341
479
  logger.info("Stopping vLLM engine")
342
480
  if self._check_health_task:
343
481
  self._check_health_task.cancel()
344
- if model_executor := getattr(self._engine.engine, "model_executor", None):
345
- model_executor.shutdown()
346
- self._engine = None
482
+ if self._engine:
483
+ if model_executor := getattr(self._engine.engine, "model_executor", None):
484
+ model_executor.shutdown()
485
+ self._engine = None
347
486
 
348
487
  async def init_xavier(self):
349
488
  await self._engine.init_xavier()
@@ -374,16 +513,18 @@ class VLLMModel(LLM):
374
513
  if model_config is None:
375
514
  model_config = VLLMModelConfig()
376
515
 
377
- cuda_count = self._get_cuda_count()
378
-
379
516
  model_config.setdefault("tokenizer_mode", "auto")
380
517
  model_config.setdefault("trust_remote_code", True)
381
- model_config.setdefault("tensor_parallel_size", cuda_count)
518
+ model_config.setdefault("tensor_parallel_size", self._device_count) # type: ignore
519
+ model_config.setdefault("pipeline_parallel_size", self._n_worker) # type: ignore
382
520
  model_config.setdefault("block_size", 16)
383
521
  model_config.setdefault("swap_space", 4)
384
522
  model_config.setdefault("gpu_memory_utilization", 0.90)
385
523
  model_config.setdefault("max_num_seqs", 256)
386
- model_config.setdefault("quantization", None)
524
+ if "model_quantization" in model_config:
525
+ model_config["quantization"] = model_config.pop("model_quantization")
526
+ else:
527
+ model_config.setdefault("quantization", None)
387
528
  model_config.setdefault("max_model_len", None)
388
529
  model_config.setdefault("guided_decoding_backend", "outlines")
389
530
  model_config.setdefault("reasoning_content", False)
@@ -0,0 +1,314 @@
1
+ # Copyright 2022-2025 XProbe Inc.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+ import asyncio
16
+ import logging
17
+ import os
18
+ from functools import partial
19
+ from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
20
+
21
+ import xoscar as xo
22
+ from vllm.executor.executor_base import DistributedExecutorBase
23
+ from vllm.utils import _run_task_with_lock, get_distributed_init_method
24
+ from vllm.worker.worker_base import WorkerWrapperBase
25
+ from xoscar.utils import get_next_port
26
+
27
+ if TYPE_CHECKING:
28
+ from vllm.config import VllmConfig
29
+ from vllm.model_executor.layers.sampler import SamplerOutput
30
+ from vllm.sequence import ExecuteModelRequest
31
+
32
+ logger = logging.getLogger(__name__)
33
+
34
+
35
+ class WorkerActor(xo.StatelessActor):
36
+ def __init__(self, vllm_config: "VllmConfig", rpc_rank: int = 0, **kwargs):
37
+ super().__init__(**kwargs)
38
+ self._worker = WorkerWrapperBase(vllm_config, rpc_rank=rpc_rank)
39
+
40
+ async def __post_create__(self):
41
+ try:
42
+ # Change process title for model
43
+ import setproctitle
44
+
45
+ setproctitle.setproctitle(f"Xinf vLLM worker: {self._worker.rpc_rank}")
46
+ except ImportError:
47
+ pass
48
+
49
+ def __getattr__(self, item):
50
+ return getattr(self._worker, item)
51
+
52
+ @classmethod
53
+ def gen_uid(cls, rank):
54
+ return f"VllmWorker_{rank}"
55
+
56
+ def execute_method(self, method: Union[str, Callable], *args, **kwargs):
57
+ logger.debug(
58
+ "Calling method %s in vllm worker %s, args: %s, kwargs: %s",
59
+ method,
60
+ self.uid,
61
+ args,
62
+ kwargs,
63
+ )
64
+ if isinstance(method, str):
65
+ return getattr(self._worker, method)(*args, **kwargs)
66
+ else:
67
+ return method(self._worker, *args, **kwargs)
68
+
69
+
70
+ class WorkerWrapper:
71
+ def __init__(
72
+ self,
73
+ loop: asyncio.AbstractEventLoop,
74
+ worker_actor_ref: xo.ActorRefType[WorkerActor],
75
+ ):
76
+ self._loop = loop
77
+ self._worker_actor_ref = worker_actor_ref
78
+
79
+ def execute_method(self, method: Union[str, Callable], *args, **kwargs):
80
+ coro = self._worker_actor_ref.execute_method(method, *args, **kwargs)
81
+ return asyncio.run_coroutine_threadsafe(coro, self._loop)
82
+
83
+ async def execute_method_async(self, method: Union[str, Callable], *args, **kwargs):
84
+ return await self._worker_actor_ref.execute_method(method, *args, **kwargs)
85
+
86
+ def kill(self):
87
+ coro = xo.kill_actor(self._worker_actor_ref)
88
+ return asyncio.run_coroutine_threadsafe(coro, self._loop)
89
+
90
+
91
+ class XinferenceDistributedExecutor(DistributedExecutorBase):
92
+ """Xoscar based distributed executor"""
93
+
94
+ use_ray: bool = False
95
+ _loop: asyncio.AbstractEventLoop
96
+ _pool_addresses: List[str]
97
+ _n_worker: int
98
+
99
+ def __init__(
100
+ self,
101
+ vllm_config: "VllmConfig",
102
+ pool_addresses: List[str],
103
+ n_worker: int,
104
+ loop: asyncio.AbstractEventLoop,
105
+ *args,
106
+ **kwargs,
107
+ ):
108
+ self._pool_addresses = pool_addresses
109
+ self._loop = loop
110
+ self._n_worker = n_worker
111
+ super().__init__(vllm_config, *args, **kwargs)
112
+
113
+ def _init_executor(self) -> None:
114
+ # Create the parallel GPU workers.
115
+ world_size = self.parallel_config.world_size
116
+ tensor_parallel_size = self.parallel_config.tensor_parallel_size
117
+
118
+ self.driver_worker: Optional[WorkerActor] = None
119
+ # The remaining workers are Xoscar actors
120
+ self.workers: List[WorkerWrapper] = []
121
+
122
+ assert (
123
+ self._pool_addresses and len(self._pool_addresses) == world_size
124
+ ), f"Pool addresses(#{len(self._pool_addresses or [])} must be equal to worldsize(#{world_size})"
125
+
126
+ futures = []
127
+ for rank in range(world_size):
128
+ coro = xo.create_actor(
129
+ WorkerActor,
130
+ self.vllm_config,
131
+ rpc_rank=rank,
132
+ address=self._pool_addresses[rank],
133
+ uid=WorkerActor.gen_uid(rank),
134
+ )
135
+ futures.append(asyncio.run_coroutine_threadsafe(coro, self._loop))
136
+ refs = [fut.result() for fut in futures]
137
+ self.workers = [WorkerWrapper(self._loop, ref) for ref in refs[1:]]
138
+ self.driver_worker = WorkerActor(self.vllm_config, rpc_rank=0)
139
+
140
+ def driver_execute_method(*args, **kwargs):
141
+ func = partial(self.driver_worker.execute_method, *args, **kwargs)
142
+ return self._loop.run_in_executor(None, func)
143
+
144
+ self.driver_exec_method = driver_execute_method
145
+
146
+ # Set environment variables for the driver and workers.
147
+ all_args_to_update_environment_variables: List[Dict[str, str]] = [
148
+ dict() for _ in range(world_size)
149
+ ]
150
+
151
+ for args in all_args_to_update_environment_variables:
152
+ # some carry-over env vars from the driver
153
+ # TODO: refactor platform-specific env vars
154
+ for name in [
155
+ "VLLM_ATTENTION_BACKEND",
156
+ "TPU_CHIPS_PER_HOST_BOUNDS",
157
+ "TPU_HOST_BOUNDS",
158
+ "VLLM_USE_V1",
159
+ "VLLM_TRACE_FUNCTION",
160
+ ]:
161
+ if name in os.environ:
162
+ args[name] = os.environ[name]
163
+
164
+ self._env_vars_for_all_workers = all_args_to_update_environment_variables
165
+
166
+ self._run_workers(
167
+ "update_environment_variables", self._env_vars_for_all_workers
168
+ )
169
+
170
+ all_kwargs = []
171
+ distributed_init_method = get_distributed_init_method(
172
+ self._pool_addresses[0].split(":", 1)[0], get_next_port()
173
+ )
174
+ for rank in range(world_size):
175
+ local_rank = rank % (world_size // self._n_worker)
176
+ kwargs = dict(
177
+ vllm_config=self.vllm_config,
178
+ local_rank=local_rank,
179
+ rank=rank,
180
+ distributed_init_method=distributed_init_method,
181
+ is_driver_worker=not self.parallel_config
182
+ or (rank % tensor_parallel_size == 0),
183
+ )
184
+ all_kwargs.append(kwargs)
185
+ self._run_workers("init_worker", all_kwargs)
186
+ self._run_workers("init_device")
187
+ self._run_workers(
188
+ "load_model",
189
+ max_concurrent_workers=self.parallel_config.max_parallel_loading_workers,
190
+ )
191
+
192
+ # This is the list of workers that are rank 0 of each TP group EXCEPT
193
+ # global rank 0. These are the workers that will broadcast to the
194
+ # rest of the workers.
195
+ self.tp_driver_workers: List[WorkerWrapper] = []
196
+ # This is the list of workers that are not drivers and not the first
197
+ # worker in a TP group. These are the workers that will be
198
+ # broadcasted to.
199
+ self.non_driver_workers: List[WorkerWrapper] = []
200
+
201
+ # Enforce rank order for correct rank to return final output.
202
+ for index, worker in enumerate(self.workers):
203
+ # The driver worker is rank 0 and not in self.workers.
204
+ rank = index + 1
205
+ if rank % self.parallel_config.tensor_parallel_size == 0:
206
+ self.tp_driver_workers.append(worker)
207
+ else:
208
+ self.non_driver_workers.append(worker)
209
+
210
+ self.pp_locks: Optional[List[asyncio.Lock]] = None
211
+
212
+ def _run_workers(
213
+ self,
214
+ method: Union[str, Callable],
215
+ *args,
216
+ async_run_tensor_parallel_workers_only: bool = False,
217
+ max_concurrent_workers: Optional[int] = None,
218
+ **kwargs,
219
+ ) -> Any:
220
+ if max_concurrent_workers:
221
+ raise NotImplementedError("max_concurrent_workers is not supported yet.")
222
+
223
+ workers = self.workers
224
+ if async_run_tensor_parallel_workers_only:
225
+ workers = self.non_driver_workers
226
+ worker_outputs = [
227
+ worker.execute_method(method, *args, **kwargs) for worker in workers
228
+ ]
229
+
230
+ if async_run_tensor_parallel_workers_only:
231
+ return worker_outputs
232
+
233
+ driver_worker_outputs = [
234
+ self.driver_worker.execute_method(method, *args, **kwargs) # type: ignore
235
+ ]
236
+ return driver_worker_outputs + [output.result() for output in worker_outputs]
237
+
238
+ def _wait_for_tasks_completion(self, parallel_worker_tasks: Any) -> None:
239
+ """Wait for futures returned from _run_workers() with
240
+ async_run_remote_workers_only to complete."""
241
+ for result in parallel_worker_tasks:
242
+ result.get()
243
+
244
+ def check_health(self) -> None:
245
+ # Assume that the workers are healthy.
246
+ # TODO: check the health by checking if the workers all alive
247
+ return
248
+
249
+ def shutdown(self) -> None:
250
+ try:
251
+ futs = [worker.kill() for worker in self.workers]
252
+ _ = [fut.result() for fut in futs]
253
+ except (RuntimeError, ConnectionError):
254
+ # event loop closed already, ignore
255
+ pass
256
+
257
+ def __del__(self):
258
+ return self.shutdown()
259
+
260
+ def _driver_execute_model(
261
+ self, execute_model_req: Optional["ExecuteModelRequest"]
262
+ ) -> Optional[List["SamplerOutput"]]:
263
+ return self.driver_worker.execute_method("execute_model", execute_model_req) # type: ignore
264
+
265
+ async def _driver_execute_model_async(
266
+ self,
267
+ execute_model_req: Optional["ExecuteModelRequest"] = None,
268
+ ) -> List["SamplerOutput"]:
269
+ if not self.tp_driver_workers:
270
+ return await self.driver_exec_method("execute_model", execute_model_req)
271
+
272
+ if self.pp_locks is None:
273
+ # This locks each pipeline parallel stage so multiple virtual
274
+ # engines can't execute on the same stage at the same time
275
+ # We create the locks here to avoid creating them in the constructor
276
+ # which uses a different asyncio loop.
277
+ self.pp_locks = [
278
+ asyncio.Lock()
279
+ for _ in range(self.parallel_config.pipeline_parallel_size)
280
+ ]
281
+
282
+ tasks = [
283
+ asyncio.create_task(
284
+ _run_task_with_lock(
285
+ self.driver_exec_method,
286
+ self.pp_locks[0],
287
+ "execute_model",
288
+ execute_model_req,
289
+ )
290
+ )
291
+ ]
292
+ for pp_rank, driver_worker in enumerate(self.tp_driver_workers, start=1):
293
+ tasks.append(
294
+ asyncio.create_task(
295
+ _run_task_with_lock(
296
+ driver_worker.execute_method_async,
297
+ self.pp_locks[pp_rank],
298
+ "execute_model",
299
+ execute_model_req,
300
+ )
301
+ )
302
+ )
303
+
304
+ results = await asyncio.gather(*tasks)
305
+
306
+ # Only the last PP stage has the final results.
307
+ return results[-1]
308
+
309
+ async def _start_worker_execution_loop(self):
310
+ coros = [
311
+ worker.execute_method_async("start_worker_execution_loop")
312
+ for worker in self.non_driver_workers
313
+ ]
314
+ return await asyncio.gather(*coros)
@@ -106,9 +106,10 @@ def generate_rerank_description(model_spec: RerankModelSpec) -> Dict[str, List[D
106
106
  return res
107
107
 
108
108
 
109
- class _ModelWrapper:
109
+ class _ModelWrapper(nn.Module):
110
110
  def __init__(self, module: nn.Module):
111
- self._module = module
111
+ super().__init__()
112
+ self.model = module
112
113
  self._local_data = threading.local()
113
114
 
114
115
  @property
@@ -116,18 +117,22 @@ class _ModelWrapper:
116
117
  return getattr(self._local_data, "n_tokens", 0)
117
118
 
118
119
  @n_tokens.setter
119
- def n_tokens(self, new_n_tokens):
120
- self._local_data.n_tokens = new_n_tokens
120
+ def n_tokens(self, value):
121
+ self._local_data.n_tokens = value
121
122
 
122
- def __getattr__(self, attr):
123
- return getattr(self._module, attr)
124
-
125
- def __call__(self, **kwargs):
126
- attention_mask = kwargs["attention_mask"]
123
+ def forward(self, **kwargs):
124
+ attention_mask = kwargs.get("attention_mask")
127
125
  # when batching, the attention mask 1 means there is a token
128
126
  # thus we just sum up it to get the total number of tokens
129
- self.n_tokens += attention_mask.sum().item()
130
- return self._module(**kwargs)
127
+ if attention_mask is not None:
128
+ self.n_tokens += attention_mask.sum().item()
129
+ return self.model(**kwargs)
130
+
131
+ def __getattr__(self, attr):
132
+ try:
133
+ return super().__getattr__(attr)
134
+ except AttributeError:
135
+ return getattr(self.model, attr)
131
136
 
132
137
 
133
138
  class RerankModel:
@@ -0,0 +1,31 @@
1
+ # Copyright (c) 2023-2024 DeepSeek.
2
+ #
3
+ # Permission is hereby granted, free of charge, to any person obtaining a copy of
4
+ # this software and associated documentation files (the "Software"), to deal in
5
+ # the Software without restriction, including without limitation the rights to
6
+ # use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
7
+ # the Software, and to permit persons to whom the Software is furnished to do so,
8
+ # subject to the following conditions:
9
+ #
10
+ # The above copyright notice and this permission notice shall be included in all
11
+ # copies or substantial portions of the Software.
12
+ #
13
+ # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
14
+ # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
15
+ # FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
16
+ # COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
17
+ # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
18
+ # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
19
+
20
+
21
+ # check if python version is above 3.10
22
+ import sys
23
+
24
+ if sys.version_info >= (3, 10):
25
+ print("Python version is above 3.10, patching the collections module.")
26
+ # Monkey patch collections
27
+ import collections
28
+ import collections.abc
29
+
30
+ for type_name in collections.abc.__all__:
31
+ setattr(collections, type_name, getattr(collections.abc, type_name))
@@ -0,0 +1,26 @@
1
+ # Copyright (c) 2023-2024 DeepSeek.
2
+ #
3
+ # Permission is hereby granted, free of charge, to any person obtaining a copy of
4
+ # this software and associated documentation files (the "Software"), to deal in
5
+ # the Software without restriction, including without limitation the rights to
6
+ # use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
7
+ # the Software, and to permit persons to whom the Software is furnished to do so,
8
+ # subject to the following conditions:
9
+ #
10
+ # The above copyright notice and this permission notice shall be included in all
11
+ # copies or substantial portions of the Software.
12
+ #
13
+ # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
14
+ # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
15
+ # FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
16
+ # COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
17
+ # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
18
+ # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
19
+
20
+ from .processing_deepseek_vl_v2 import DeepseekVLV2Processor
21
+ from .modeling_deepseek_vl_v2 import DeepseekVLV2ForCausalLM
22
+
23
+ __all__ = [
24
+ "DeepseekVLV2Processor",
25
+ "DeepseekVLV2ForCausalLM",
26
+ ]