arbor-ai 0.1.8__py3-none-any.whl → 0.1.9__py3-none-any.whl

arbor/server/services/comms/comms.py

@@ -49,10 +49,15 @@ class ArborServerCommsHandler:
             yield status
 
     def close(self):
+        self.command_socket.setsockopt(zmq.LINGER, 0)
         self.command_socket.close()
+        self.status_socket.setsockopt(zmq.LINGER, 0)
         self.status_socket.close()
+        self.data_socket.setsockopt(zmq.LINGER, 0)
         self.data_socket.close()
+        self.broadcast_socket.setsockopt(zmq.LINGER, 0)
         self.broadcast_socket.close()
+        self.handshake_socket.setsockopt(zmq.LINGER, 0)
         self.handshake_socket.close()
         self.context.term()
 

arbor/server/services/grpo_manager.py

@@ -276,7 +276,9 @@ class GRPOManager:
                 inference_manager.kill()
 
             # Send termination command through REQ socket
-            self.server_comms_handler.send_broadcast({"message": "terminate"})
+            # self.server_comms_handler.send_broadcast({"message": "terminate"})
+            self.training_process.terminate()
+            print("Waiting for training process to finish")
 
             # Wait for training process to finish
             if self.training_process:
@@ -289,6 +291,21 @@ class GRPOManager:
             if self.server_comms_handler:
                 self.server_comms_handler.close()
 
+            # Force kill training process if still running
+            if self.training_process and self.training_process.poll() is None:
+                self.training_process.kill()
+                self.training_process.wait()
+
+            # Reinitialize incase we want to start a new training run
+            self.training_process = None
+            self.current_model = None
+            self.server_comms_handler = None
+            self.status_thread = None
+            self.model_saved_and_reload_requested = False
+
+            self.data_count = 0
+            self.last_inference_update = 0
+
             if self.train_kwargs and "output_dir" in self.train_kwargs:
                 print(
                     f"Training completed. Model saved to {self.train_kwargs['output_dir']}"
@@ -297,9 +314,12 @@ class GRPOManager:
                     print(
                         f"Warning: Output directory {self.train_kwargs['output_dir']} does not exist"
                     )
-                return self.train_kwargs["output_dir"]
+                output_dir = self.train_kwargs["output_dir"]
+                self.train_kwargs = None
+                return output_dir
             else:
                 print("Training terminated, no output directory specified")
+                self.train_kwargs = None
                 return None
 
     def _should_update_model(self):

arbor/server/services/inference_manager.py

@@ -61,7 +61,7 @@ class InferenceManager:
         my_env = os.environ.copy()
         my_env["CUDA_VISIBLE_DEVICES"] = self.settings.arbor_config.inference.gpu_ids
         n_gpus = self.settings.arbor_config.inference.gpu_ids.count(",") + 1
-        # command = f"vllm serve {model} --port {port} --gpu-memory-utilization 0.9 --tensor-parallel-size {n_gpus} --max_model_len 8192 --enable_prefix_caching --guided-decoding-backend xgrammar"
+        # command = f"vllm serve {model} --port {port} --gpu-memory-utilization 0.9 --tensor-parallel-size {n_gpus} --max_model_len 8192 --enable_prefix_caching"
         command = f"python -m sglang_router.launch_server --model-path {model} --dp-size {n_gpus} --router-policy round_robin --port {port} --host 0.0.0.0"
         print(f"Running command: {command}")
 
@@ -137,8 +137,6 @@ class InferenceManager:
         process = self.process
         thread = self.thread
 
-        terminate_process(process)
-
         # Clear references first
         self.process = None
         self.thread = None

{arbor_ai-0.1.8.dist-info → arbor_ai-0.1.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arbor-ai
-Version: 0.1.8
+Version: 0.1.9
 Summary: A framework for fine-tuning and managing language models
 Author-email: Noah Ziems <nziems2@nd.edu>
 Project-URL: Homepage, https://github.com/Ziems/arbor

{arbor_ai-0.1.8.dist-info → arbor_ai-0.1.9.dist-info}/RECORD

@@ -17,18 +17,18 @@ arbor/server/core/logging.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,
 arbor/server/services/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 arbor/server/services/dependencies.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 arbor/server/services/file_manager.py,sha256=Z9z4A4EzvPauid_DBfpim401DDtuJy_TbX4twTWDJWI,12119
-arbor/server/services/grpo_manager.py,sha256=T-f1TrNSi_kmxPOcpaDphS8Xf3UMUbricocc6fuaKIM,12077
-arbor/server/services/inference_manager.py,sha256=qR9xPiYs4Is24vgeF72w7Hbe8j_PGEbl-qewcvUV-dA,9731
+arbor/server/services/grpo_manager.py,sha256=Nd7T0q1RTQUhhG4uMB8lUonw_6Rww31yVJo8MM-noU8,12903
+arbor/server/services/inference_manager.py,sha256=gHI-Biy3TtGkyWxIDKY-uqZZm_fiQJLktkPY8ezRvo8,9660
 arbor/server/services/job_manager.py,sha256=m_d4UPwN_82f7t7K443DaFpFoyv7JZSZKml8tawt1Bk,2186
 arbor/server/services/training_manager.py,sha256=oQdhpfxdgp_lCTb_lxhvjupdLrcg6HL3TEbct_q9F6I,21065
 arbor/server/services/comms/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-arbor/server/services/comms/comms.py,sha256=Dg08D2Fm5TAEiGyr0Qcr0uocabQpFD_sBVhxIkj9D2M,7424
+arbor/server/services/comms/comms.py,sha256=3KN3mzwPvfW2_L5hq02JdAk6yOMyhY0_pBz-DDr5A3o,7694
 arbor/server/services/scripts/grpo_training.py,sha256=V36pCMZDJj2DdzquxScOddi9zP8EVPGWN3HGiftFfrY,21082
 arbor/server/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 arbor/server/utils/helpers.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-arbor_ai-0.1.8.dist-info/licenses/LICENSE,sha256=5vFGrbOFeXXM83JV9o16w7ohH4WLeu3-57GocJSz8ow,1067
-arbor_ai-0.1.8.dist-info/METADATA,sha256=kAZBj176hfqSrvrcWb0Wz8_vU33yiZJ-ck9buyDF6Jg,2234
-arbor_ai-0.1.8.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
-arbor_ai-0.1.8.dist-info/entry_points.txt,sha256=PGBX-MfNwfIl8UPFgsX3gjtXLqSogRhOktKMpZUysD0,40
-arbor_ai-0.1.8.dist-info/top_level.txt,sha256=jzWdp3BRYqvZDMFsPajrcftvvlluzVDErkD8IMRfhYs,6
-arbor_ai-0.1.8.dist-info/RECORD,,
+arbor_ai-0.1.9.dist-info/licenses/LICENSE,sha256=5vFGrbOFeXXM83JV9o16w7ohH4WLeu3-57GocJSz8ow,1067
+arbor_ai-0.1.9.dist-info/METADATA,sha256=GVTgqz290CXbwW5Yyt-EpelGmp1UZuK8kpN67Fp2bV0,2234
+arbor_ai-0.1.9.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
+arbor_ai-0.1.9.dist-info/entry_points.txt,sha256=PGBX-MfNwfIl8UPFgsX3gjtXLqSogRhOktKMpZUysD0,40
+arbor_ai-0.1.9.dist-info/top_level.txt,sha256=jzWdp3BRYqvZDMFsPajrcftvvlluzVDErkD8IMRfhYs,6
+arbor_ai-0.1.9.dist-info/RECORD,,