update handler

Author: Yunnglin

src/twinkle/server/model/app.py

@@ -118,10 +118,10 @@ def __del__(self):
             pass
 
     async def _cleanup_adapter(self, adapter_name: str) -> None:
-        if self.get_adapter_info(adapter_name):
-            self.clear_adapter_state(adapter_name)
+        if self.get_resource_info(adapter_name):
+            self.clear_resource_state(adapter_name)
             self.model.remove_adapter(adapter_name)
-            self.unregister_adapter(adapter_name)
+            self.unregister_resource(adapter_name)
             await self.state.unload_model(adapter_name)
 
     async def _on_adapter_expired(self, adapter_name: str) -> None:

src/twinkle/server/model/tinker_handlers.py

@@ -44,12 +44,12 @@ async def _create_adapter():
                 if body.lora_config:
                     lora_cfg = LoraConfig(r=body.lora_config.rank, target_modules='all-linear')
                     adapter_name = self.get_adapter_name(adapter_name=_model_id)
-                    self.register_adapter(adapter_name, token, session_id=body.session_id)
+                    self.register_resource(adapter_name, token, session_id=body.session_id)
                     self.model.add_adapter_to_model(adapter_name=adapter_name, config_or_dir=lora_cfg)
                     self.model.set_template('Template', adapter_name=adapter_name, model_id=self.base_model)
                     self.model.set_processor('InputProcessor', adapter_name=adapter_name)
                     self.model.set_optimizer('Adam', adapter_name=adapter_name)
-                    self.set_adapter_state(adapter_name, 'grad_ready', False)
+                    self.set_resource_state(adapter_name, 'grad_ready', False)
                 training_run_manager = create_training_run_manager(token, client_type='tinker')
                 training_run_manager.save(_model_id, body)
                 return types.CreateModelResponse(model_id=_model_id)
@@ -108,7 +108,7 @@ async def forward(request: Request, body: types.ForwardRequest,
         async def _do_forward():
             try:
                 adapter_name = self.get_adapter_name(adapter_name=body.model_id)
-                self.assert_adapter_exists(adapter_name=adapter_name)
+                self.assert_resource_exists(adapter_name)
                 datum_list = body.forward_input.data
                 loss_fn_config = body.forward_input.loss_fn_config or {}
                 output = self.model.tinker_forward_only(inputs=datum_list, adapter_name=adapter_name)
@@ -149,15 +149,15 @@ async def forward_backward(
         async def _do_forward_backward():
             try:
                 adapter_name = self.get_adapter_name(adapter_name=body.model_id)
-                self.assert_adapter_exists(adapter_name=adapter_name)
+                self.assert_resource_exists(adapter_name)
                 datum_list = body.forward_backward_input.data
                 loss_fn = body.forward_backward_input.loss_fn
                 loss_fn_config = body.forward_backward_input.loss_fn_config or {}
                 output, loss = self.model.tinker_forward_backward(
                     inputs=datum_list, adapter_name=adapter_name, loss_fn=loss_fn, **loss_fn_config)
                 output_type = ('ImportanceSamplingLossReturn'
                                if loss_fn == 'importance_sampling' else 'CrossEntropyLossReturn')
-                self.set_adapter_state(adapter_name, 'grad_ready', True)
+                self.set_resource_state(adapter_name, 'grad_ready', True)
                 return types.ForwardBackwardOutput(
                     loss_fn_output_type=output_type,
                     loss_fn_outputs=output,
@@ -194,12 +194,12 @@ async def optim_step(
         async def _do_optim():
             try:
                 adapter_name = self.get_adapter_name(adapter_name=body.model_id)
-                self.assert_adapter_exists(adapter_name=adapter_name)
-                if not self.get_adapter_state(adapter_name, 'grad_ready', False):
+                self.assert_resource_exists(adapter_name)
+                if not self.get_resource_state(adapter_name, 'grad_ready', False):
                     raise RuntimeError(f'No accumulated gradients for adapter={adapter_name}; '
                                        'call forward_backward before optim_step')
                 self.model.tinker_step(adam_params=body.adam_params, adapter_name=adapter_name)
-                self.set_adapter_state(adapter_name, 'grad_ready', False)
+                self.set_resource_state(adapter_name, 'grad_ready', False)
                 metrics = self.model.tinker_calculate_metric(is_training=True, adapter_name=adapter_name)
                 return types.OptimStepResponse(metrics=metrics)
             except Exception:
@@ -222,7 +222,7 @@ async def save_weights(
         async def _do_save():
             try:
                 adapter_name = self.get_adapter_name(adapter_name=body.model_id)
-                self.assert_adapter_exists(adapter_name=adapter_name)
+                self.assert_resource_exists(adapter_name)
                 checkpoint_manager = create_checkpoint_manager(token, client_type='tinker')
                 checkpoint_name = checkpoint_manager.get_ckpt_name(body.path)
                 save_dir = checkpoint_manager.get_save_dir(model_id=body.model_id, is_sampler=False)
@@ -250,7 +250,7 @@ async def save_weights_for_sampler(
         async def _do_save_for_sampler():
             try:
                 adapter_name = self.get_adapter_name(adapter_name=body.model_id)
-                self.assert_adapter_exists(adapter_name=adapter_name)
+                self.assert_resource_exists(adapter_name)
                 checkpoint_manager = create_checkpoint_manager(token, client_type='tinker')
                 checkpoint_name = checkpoint_manager.get_ckpt_name(body.path)
                 save_dir = checkpoint_manager.get_save_dir(model_id=body.model_id, is_sampler=True)
@@ -287,10 +287,10 @@ async def _do_load():
             try:
                 assert self.model is not None, 'Model not loaded, please load model first'
                 adapter_name = self.get_adapter_name(adapter_name=body.model_id)
-                self.assert_adapter_exists(adapter_name=adapter_name)
+                self.assert_resource_exists(adapter_name)
                 self.model.tinker_load(
                     checkpoint_dir=body.path, load_optimizer=body.optimizer, adapter_name=adapter_name, token=token)
-                self.set_adapter_state(adapter_name, 'grad_ready', False)
+                self.set_resource_state(adapter_name, 'grad_ready', False)
                 return types.LoadWeightsResponse(path=body.path, type='load_weights')
             except Exception:
                 logger.error(traceback.format_exc())

src/twinkle/server/model/twinkle_handlers.py

@@ -85,7 +85,7 @@ async def forward(request: Request, body: types.ForwardRequest,
         adapter_name = _get_twinkle_adapter_name(request, body.adapter_name)
 
         async def _task():
-            self.assert_adapter_exists(adapter_name=adapter_name)
+            self.assert_resource_exists(adapter_name)
             extra_kwargs = body.model_extra or {}
             inputs = _parse_inputs(body.inputs)
             ret = self.model.forward(inputs=inputs, adapter_name=adapter_name, **extra_kwargs)
@@ -103,7 +103,7 @@ async def forward_only(
         adapter_name = _get_twinkle_adapter_name(request, body.adapter_name)
 
         async def _task():
-            self.assert_adapter_exists(adapter_name=adapter_name)
+            self.assert_resource_exists(adapter_name)
             extra_kwargs = body.model_extra or {}
             inputs = _parse_inputs(body.inputs)
             ret = self.model.forward_only(inputs=inputs, adapter_name=adapter_name, **extra_kwargs)
@@ -121,7 +121,7 @@ async def calculate_loss(
         adapter_name = _get_twinkle_adapter_name(request, body.adapter_name)
 
         async def _task():
-            self.assert_adapter_exists(adapter_name=adapter_name)
+            self.assert_resource_exists(adapter_name)
             extra_kwargs = body.model_extra or {}
             ret = self.model.calculate_loss(adapter_name=adapter_name, **extra_kwargs)
             return {'result': ret}
@@ -134,7 +134,7 @@ async def backward(request: Request, body: types.AdapterRequest, self: ModelMana
         adapter_name = _get_twinkle_adapter_name(request, body.adapter_name)
 
         async def _task():
-            self.assert_adapter_exists(adapter_name=adapter_name)
+            self.assert_resource_exists(adapter_name)
             extra_kwargs = body.model_extra or {}
             self.model.backward(adapter_name=adapter_name, **extra_kwargs)
 
@@ -157,7 +157,7 @@ def first_element(data):
             return data
 
         async def _task():
-            self.assert_adapter_exists(adapter_name=adapter_name)
+            self.assert_resource_exists(adapter_name)
             extra_kwargs = body.model_extra or {}
             all_inputs = _parse_inputs(body.inputs)
             for inputs in all_inputs:
@@ -179,7 +179,7 @@ async def clip_grad_norm(
         adapter_name = _get_twinkle_adapter_name(request, body.adapter_name)
 
         async def _task():
-            self.assert_adapter_exists(adapter_name=adapter_name)
+            self.assert_resource_exists(adapter_name)
             extra_kwargs = body.model_extra or {}
             ret = self.model.clip_grad_norm(adapter_name=adapter_name, **extra_kwargs)
             return {'result': str(ret)}
@@ -192,7 +192,7 @@ async def step(request: Request, body: types.AdapterRequest, self: ModelManageme
         adapter_name = _get_twinkle_adapter_name(request, body.adapter_name)
 
         async def _task():
-            self.assert_adapter_exists(adapter_name=adapter_name)
+            self.assert_resource_exists(adapter_name)
             extra_kwargs = body.model_extra or {}
             self.model.step(adapter_name=adapter_name, **extra_kwargs)
 
@@ -204,7 +204,7 @@ async def zero_grad(request: Request, body: types.AdapterRequest, self: ModelMan
         adapter_name = _get_twinkle_adapter_name(request, body.adapter_name)
 
         async def _task():
-            self.assert_adapter_exists(adapter_name=adapter_name)
+            self.assert_resource_exists(adapter_name)
             extra_kwargs = body.model_extra or {}
             self.model.zero_grad(adapter_name=adapter_name, **extra_kwargs)
 
@@ -216,7 +216,7 @@ async def lr_step(request: Request, body: types.AdapterRequest, self: ModelManag
         adapter_name = _get_twinkle_adapter_name(request, body.adapter_name)
 
         async def _task():
-            self.assert_adapter_exists(adapter_name=adapter_name)
+            self.assert_resource_exists(adapter_name)
             extra_kwargs = body.model_extra or {}
             self.model.lr_step(adapter_name=adapter_name, **extra_kwargs)
 
@@ -232,7 +232,7 @@ async def clip_grad_and_step(
         adapter_name = _get_twinkle_adapter_name(request, body.adapter_name)
 
         async def _task():
-            self.assert_adapter_exists(adapter_name=adapter_name)
+            self.assert_resource_exists(adapter_name)
             extra_kwargs = body.model_extra or {}
             self.model.clip_grad_and_step(
                 max_grad_norm=body.max_grad_norm,
@@ -253,7 +253,7 @@ async def get_train_configs(
         adapter_name = _get_twinkle_adapter_name(request, body.adapter_name)
 
         async def _task():
-            self.assert_adapter_exists(adapter_name=adapter_name)
+            self.assert_resource_exists(adapter_name)
             extra_kwargs = body.model_extra or {}
             ret = self.model.get_train_configs(adapter_name=adapter_name, **extra_kwargs)
             return {'result': ret}
@@ -266,7 +266,7 @@ async def set_loss(request: Request, body: types.SetLossRequest, self: ModelMana
         adapter_name = _get_twinkle_adapter_name(request, body.adapter_name)
 
         async def _task():
-            self.assert_adapter_exists(adapter_name=adapter_name)
+            self.assert_resource_exists(adapter_name)
             extra_kwargs = body.model_extra or {}
             self.model.set_loss(body.loss_cls, adapter_name=adapter_name, **extra_kwargs)
 
@@ -282,7 +282,7 @@ async def set_optimizer(
         adapter_name = _get_twinkle_adapter_name(request, body.adapter_name)
 
         async def _task():
-            self.assert_adapter_exists(adapter_name=adapter_name)
+            self.assert_resource_exists(adapter_name)
             extra_kwargs = body.model_extra or {}
             self.model.set_optimizer(body.optimizer_cls, adapter_name=adapter_name, **extra_kwargs)
 
@@ -298,7 +298,7 @@ async def set_lr_scheduler(
         adapter_name = _get_twinkle_adapter_name(request, body.adapter_name)
 
         async def _task():
-            self.assert_adapter_exists(adapter_name=adapter_name)
+            self.assert_resource_exists(adapter_name)
             extra_kwargs = body.model_extra or {}
             self.model.set_lr_scheduler(body.scheduler_cls, adapter_name=adapter_name, **extra_kwargs)
 
@@ -311,7 +311,7 @@ async def save(request: Request, body: types.SaveRequest,
         adapter_name = _get_twinkle_adapter_name(request, body.adapter_name)
 
         async def _task():
-            self.assert_adapter_exists(adapter_name=adapter_name)
+            self.assert_resource_exists(adapter_name)
             extra_kwargs = body.model_extra or {}
             checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle')
             checkpoint_name = checkpoint_manager.get_ckpt_name(body.name)
@@ -333,7 +333,7 @@ async def load(request: Request, body: types.LoadRequest, self: ModelManagement
         adapter_name = _get_twinkle_adapter_name(request, body.adapter_name)
 
         async def _task():
-            self.assert_adapter_exists(adapter_name=adapter_name)
+            self.assert_resource_exists(adapter_name)
             extra_kwargs = body.model_extra or {}
             checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle')
             resolved = checkpoint_manager.resolve_load_path(body.name)
@@ -393,7 +393,7 @@ async def _task():
             config = deserialize_object(body.config)
             extra_kwargs = body.model_extra or {}
             training_run_manager = create_training_run_manager(token, client_type='twinkle')
-            self.register_adapter(adapter_name, token, session_id=session_id)
+            self.register_resource(adapter_name, token, session_id)
             self.model.add_adapter_to_model(adapter_name, config, **extra_kwargs)
 
             lora_config = None
@@ -416,7 +416,7 @@ async def apply_patch(
         adapter_name = _get_twinkle_adapter_name(request, body.adapter_name)
 
         async def _task():
-            self.assert_adapter_exists(adapter_name=adapter_name)
+            self.assert_resource_exists(adapter_name)
             extra_kwargs = body.model_extra or {}
             patch_cls = deserialize_object(body.patch_cls)
             self.model.apply_patch(patch_cls, adapter_name=adapter_name, **extra_kwargs)
@@ -433,7 +433,7 @@ async def add_metric(
         adapter_name = _get_twinkle_adapter_name(request, body.adapter_name)
 
         async def _task():
-            self.assert_adapter_exists(adapter_name=adapter_name)
+            self.assert_resource_exists(adapter_name)
             extra_kwargs = body.model_extra or {}
             metric_cls = deserialize_object(body.metric_cls)
             self.model.add_metric(metric_cls, is_training=body.is_training, adapter_name=adapter_name, **extra_kwargs)
@@ -450,7 +450,7 @@ async def set_template(
         adapter_name = _get_twinkle_adapter_name(request, body.adapter_name)
 
         async def _task():
-            self.assert_adapter_exists(adapter_name=adapter_name)
+            self.assert_resource_exists(adapter_name)
             extra_kwargs = body.model_extra or {}
             self.model.set_template(body.template_cls, adapter_name=adapter_name, **extra_kwargs)
 
@@ -466,7 +466,7 @@ async def set_processor(
         adapter_name = _get_twinkle_adapter_name(request, body.adapter_name)
 
         async def _task():
-            self.assert_adapter_exists(adapter_name=adapter_name)
+            self.assert_resource_exists(adapter_name)
             extra_kwargs = body.model_extra or {}
             self.model.set_processor(body.processor_cls, adapter_name=adapter_name, **extra_kwargs)
 
@@ -482,7 +482,7 @@ async def calculate_metric(
         adapter_name = _get_twinkle_adapter_name(request, body.adapter_name)
 
         async def _task():
-            self.assert_adapter_exists(adapter_name=adapter_name)
+            self.assert_resource_exists(adapter_name)
             extra_kwargs = body.model_extra or {}
             ret = self.model.calculate_metric(is_training=body.is_training, adapter_name=adapter_name, **extra_kwargs)
             return {'result': ret}
@@ -499,7 +499,7 @@ async def get_state_dict(
         adapter_name = _get_twinkle_adapter_name(request, body.adapter_name)
 
         async def _task():
-            self.assert_adapter_exists(adapter_name=adapter_name)
+            self.assert_resource_exists(adapter_name)
             extra_kwargs = body.model_extra or {}
             ret = self.model.get_state_dict(adapter_name=adapter_name, **extra_kwargs)
             return {'result': ret}

src/twinkle/server/processor/app.py

@@ -84,7 +84,7 @@ async def _ensure_sticky(self):
     def _on_processor_expired(self, processor_id: str) -> None:
         """Called by the countdown thread when a processor's session expires."""
         self.resource_dict.pop(processor_id, None)
-        self.unregister_processor(processor_id)
+        self.unregister_resource(processor_id)
 
 
 def build_processor_app(ncpu_proc_per_node: int,

src/twinkle/server/processor/twinkle_handlers.py

@@ -53,7 +53,7 @@ async def create(
         processor_id = str(uuid.uuid4().hex)
 
         # Register for lifecycle tracking (enforces per-user limit)
-        self.register_processor(processor_id, token, session_id)
+        self.register_resource(processor_id, token, session_id)
 
         _kwargs.pop('remote_group', None)
         _kwargs.pop('device_mesh', None)
@@ -91,7 +91,7 @@ async def call(
         function_name = body.function
         _kwargs = body.model_extra or {}
         processor_id = processor_id[4:]
-        self.assert_processor_exists(processor_id=processor_id)
+        self.assert_resource_exists(processor_id)
         processor = self.resource_dict.get(processor_id)
         function = getattr(processor, function_name, None)