mantisdk 0.1.0__py3-none-any.whl

mantisdk/utils/system_snapshot.py

@@ -0,0 +1,81 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import platform
+import socket
+from contextlib import suppress
+from datetime import datetime
+from typing import Any, Dict, List, cast
+
+import psutil
+from gpustat import GPUStat, GPUStatCollection
+
+
+def system_snapshot(include_gpu: bool = False) -> Dict[str, Any]:
+    """Capture a snapshot of the system's hardware and software information.
+
+    Args:
+        include_gpu: Whether to include GPU information.
+
+    Returns:
+        A dictionary containing the system's hardware and software information.
+    """
+    # CPU
+    cpu = {
+        "cpu_name": platform.processor(),
+        "cpu_cores": psutil.cpu_count(logical=False),
+        "cpu_threads": psutil.cpu_count(logical=True),
+        "cpu_usage_pct": psutil.cpu_percent(0.0),
+    }
+
+    # Memory
+    vm = psutil.virtual_memory()
+    mem = {
+        "mem_used_gb": round(vm.used / (2**30), 2),
+        "mem_total_gb": round(vm.total / (2**30), 2),
+        "mem_pct": vm.percent,
+    }
+
+    # Disk
+    du = psutil.disk_usage("/")
+    disk = {
+        "disk_used_gb": round(du.used / (2**30), 2),
+        "disk_total_gb": round(du.total / (2**30), 2),
+        "disk_pct": du.percent,
+    }
+
+    # GPU (only query if explicitly requested)
+    gpus: List[Dict[str, Any]] = []
+    if include_gpu:
+        with suppress(Exception):
+            for g in GPUStatCollection.new_query().gpus:  # type: ignore
+                g = cast(GPUStat, g)
+                gpus.append(
+                    {
+                        "gpu": g.name,  # type: ignore
+                        "util_pct": g.utilization,
+                        "mem_used_mb": g.memory_used,
+                        "mem_total_mb": g.memory_total,
+                        "temp_c": g.temperature,
+                    }
+                )
+
+    # Network
+    net = psutil.net_io_counters()
+    netinfo = {
+        "bytes_sent_mb": round(net.bytes_sent / (2**20), 2),
+        "bytes_recv_mb": round(net.bytes_recv / (2**20), 2),
+    }
+
+    # OS / meta
+    return {
+        "timestamp": datetime.now().isoformat(timespec="seconds"),
+        "host": socket.gethostname(),
+        "os": platform.platform(),
+        **cpu,
+        **mem,
+        **disk,
+        **netinfo,
+        **({"gpus": gpus} if include_gpu else {}),
+    }

mantisdk/verl/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""This package contains a *hacky* integration of VERL with Mantisdk."""
+
+from .daemon import *
+from .dataset import *
+from .entrypoint import *
+from .trainer import *

mantisdk/verl/__main__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .entrypoint import main
+
+if __name__ == "__main__":
+    main()

mantisdk/verl/async_server.py

@@ -0,0 +1,46 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# type: ignore
+
+from copy import deepcopy
+
+import ray
+from starlette.requests import Request
+from starlette.responses import JSONResponse, StreamingResponse
+from verl.workers.rollout.vllm_rollout.vllm_async_server import AsyncvLLMServer
+from vllm.entrypoints.openai.protocol import ChatCompletionRequest, ErrorResponse
+
+from mantisdk.instrumentation.vllm import ChatCompletionResponsePatched, instrument_vllm
+
+
+def _unwrap_ray_remote(cls):
+    if hasattr(cls, "__ray_actor_class__"):
+        cls = cls.__ray_actor_class__
+    return cls
+
+
+@ray.remote(num_cpus=1)
+class PatchedvLLMServer(_unwrap_ray_remote(AsyncvLLMServer)):
+
+    def __init__(self, *args, **kwargs):
+        instrument_vllm()
+        super().__init__(*args, **kwargs)
+
+        self.config = deepcopy(self.config)
+        self.config.rollout.multi_turn.tool_config_path = "/dev/null"
+
+    async def chat_completion(self, raw_request: Request):
+        """OpenAI-compatible HTTP endpoint.
+
+        API reference: [OpenAI-compatible server documentation](https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html)
+        """
+        request_json = await raw_request.json()
+        request = ChatCompletionRequest(**request_json)
+        generator = await self.openai_serving_chat.create_chat_completion(request, raw_request)
+
+        if isinstance(generator, ErrorResponse):
+            return JSONResponse(content=generator.model_dump(), status_code=generator.code)
+        if request.stream:
+            return StreamingResponse(content=generator, media_type="text/event-stream")
+        else:
+            return JSONResponse(content=generator.model_dump())

mantisdk/verl/config.yaml

@@ -0,0 +1,27 @@
+hydra:
+  searchpath:
+    - pkg://verl/trainer/config
+
+defaults:
+  - ppo_trainer
+  - _self_
+
+mantisdk:
+  port: 9999
+  trace_aggregator:
+    level: transition  # transition or trajectory, docs refer to https://mantisdk.github.io/posts/trajectory_level_aggregation/
+    trajectory_max_prompt_length: 2048  # supported in trajectory level aggregation, suggest to set as maximum length for the prompt in first turn
+    trajectory_max_response_length: 8192  # supported in trajectory level aggregation, suggest to set as maximum length for the cumulative agent responses in the full trajectory, i.e., n_turns * (max_response_length + max_prompt_length)
+    debug: False  # supported in trajectory level aggregation, enable to diagnose trace merging failures
+    mismatch_log_dir: ./mismatch_cases  # supported in trajectory level aggregation with debug=True, directory to store logs of mismatch cases
+
+data:
+  filter_overlong_prompts: false
+
+actor_rollout_ref:
+  rollout:
+    mode: async
+    agent:
+      custom_async_server:
+        path: pkg://mantisdk.verl.async_server
+        name: PatchedvLLMServer