wip

Author: tastelikefeet

src/twinkle/infra/_ray/resource_manager.py

@@ -137,6 +137,19 @@ def __init__(self, nproc_per_node: int, ncpu_proc_per_node: int, groups: List[De
         if self.node_ranks.count(0) > 1:
             self.node_ranks = list(range(len(self.placement_groups)))
 
+        self.visible_devices = []
+
+        @ray.remote
+        def get_visible_devices():
+            return os.environ.get(Platform.get_platform(group.device_type).visible_device_env())
+
+        if self.placement_groups:
+            self.visible_devices = ray.get([
+                get_visible_devices.options(placement_group=pg).remote() for pg in self.placement_groups
+            ])
+
+        breakpoint()
+
         self.node2pg: Dict[int, PlacementGroup] = {}
         # Map actual node indices to placement groups
         # For GPU/NPU groups, node indices start from self.min_node_idx

src/twinkle/server/tinker/model.py

@@ -100,15 +100,17 @@ def __init__(self,
             else:
                 self.device_mesh = DeviceMesh.from_sizes(**device_mesh)
             self.use_megatron = use_megatron
+            replica_context = serve.get_replica_context()
+            replica_id = replica_context.replica_id.unique_id
             # Initialize model immediately - choose backend based on use_megatron
             if use_megatron:
                 from .common.megatron_model import TwinkleCompatMegatronModel
                 self.model = TwinkleCompatMegatronModel(
-                    model_id=model_id, device_mesh=self.device_mesh, remote_group=self.device_group.name, **kwargs)
+                    model_id=model_id, device_mesh=self.device_mesh, remote_group=self.device_group.name, instance_id=replica_id, **kwargs)
             else:
                 from .common.transformers_model import TwinkleCompatTransformersModel
                 self.model = TwinkleCompatTransformersModel(
-                    model_id=model_id, device_mesh=self.device_mesh, remote_group=self.device_group.name, **kwargs)
+                    model_id=model_id, device_mesh=self.device_mesh, remote_group=self.device_group.name, instance_id=replica_id, **kwargs)
             self.base_model = model_id
             self.state: ServerStateProxy = get_server_state()
 

src/twinkle/server/tinker/sampler.py

@@ -102,6 +102,8 @@ def __init__(self,
             else:
                 self.device_mesh = DeviceMesh.from_sizes(**device_mesh)
             self.sampler_type = sampler_type
+            replica_context = serve.get_replica_context()
+            replica_id = replica_context.replica_id.unique_id
 
             # Initialize sampler based on type
             if sampler_type == 'vllm':
@@ -112,6 +114,7 @@ def __init__(self,
                     engine_args=sampler_kwargs,
                     device_mesh=self.device_mesh,
                     remote_group=self.device_group.name,
+                    instance_id=replica_id,
                     **{
                         k: v
                         for k, v in kwargs.items() if k not in ['engine_args']

src/twinkle/server/twinkle/model.py

@@ -171,14 +171,16 @@ def __init__(self, nproc_per_node: int, device_group: Dict[str, Any], device_mes
                 self.device_mesh = DeviceMesh(**device_mesh)
             else:
                 self.device_mesh = DeviceMesh.from_sizes(**device_mesh)
+            replica_context = serve.get_replica_context()
+            replica_id = replica_context.replica_id.unique_id
             if use_megatron:
                 from twinkle.model import MultiLoraMegatronModel
                 self.model = MultiLoraMegatronModel(
-                    model_id=model_id, device_mesh=self.device_mesh, remote_group=self.device_group.name, **kwargs)
+                    model_id=model_id, device_mesh=self.device_mesh, remote_group=self.device_group.name, instance_id=replica_id, **kwargs)
             else:
                 from twinkle.model import MultiLoraTransformersModel
                 self.model = MultiLoraTransformersModel(
-                    model_id=model_id, device_mesh=self.device_mesh, remote_group=self.device_group.name, **kwargs)
+                    model_id=model_id, device_mesh=self.device_mesh, remote_group=self.device_group.name, instance_id=replica_id, **kwargs)
 
             # Initialize state before adapter manager (mixin needs self.state)
             self.state: ServerStateProxy = get_server_state()

src/twinkle/server/twinkle/sampler.py

@@ -152,7 +152,8 @@ def __init__(self,
             else:
                 self.device_mesh = DeviceMesh.from_sizes(**device_mesh)
             self.sampler_type = sampler_type
-
+            replica_context = serve.get_replica_context()
+            replica_id = replica_context.replica_id.unique_id
             # Initialize sampler based on type
             if sampler_type == 'vllm':
                 from twinkle.sampler import vLLMSampler
@@ -162,14 +163,15 @@ def __init__(self,
                     engine_args=sampler_kwargs,
                     device_mesh=self.device_mesh,
                     remote_group=self.device_group.name,
+                    instance_id=replica_id,
                     **{
                         k: v
                         for k, v in kwargs.items() if k not in ['engine_args']
                     })
             else:
                 from twinkle.sampler import TorchSampler
                 self.sampler = TorchSampler(
-                    model_id=model_id, device_mesh=self.device_mesh, remote_group=self.device_group.name, **kwargs)
+                    model_id=model_id, device_mesh=self.device_mesh, instance_id=replica_id, remote_group=self.device_group.name, **kwargs)
 
             # Initialize state and adapter manager
             self.state: ServerStateProxy = get_server_state()