mle-kit-mcp 0.0.2__py3-none-any.whl → 0.0.4__py3-none-any.whl

mle_kit_mcp/files.py

@@ -1,12 +1,25 @@
+import os
+from typing import Optional
 from pathlib import Path
 
 DIR_PATH = Path(__file__).parent
 ROOT_PATH = DIR_PATH.parent
+DEFAULT_WORKSPACE_DIR_PATH: Path = DIR_PATH / "workdir"
 
-WORKSPACE_DIR_PATH: Path = ROOT_PATH / "workdir"
 
+class WorkspaceDirectory:
+    workspace_dir: Optional[Path] = None
 
-def set_workspace_dir(path: Path) -> None:
-    global WORKSPACE_DIR_PATH
-    path.mkdir(parents=True, exist_ok=True)
-    WORKSPACE_DIR_PATH = path
+    @classmethod
+    def get_dir(cls) -> Path:
+        if cls.workspace_dir is None:
+            return Path(os.getenv("WORKSPACE_DIR", DEFAULT_WORKSPACE_DIR_PATH))
+        return cls.workspace_dir
+
+    @classmethod
+    def set_dir(cls, workspace_dir: Path) -> None:
+        cls.workspace_dir = workspace_dir
+
+
+def get_workspace_dir() -> Path:
+    return WorkspaceDirectory.get_dir()

mle_kit_mcp/server.py

@@ -1,9 +1,10 @@
-import fire  # type: ignore
 from pathlib import Path
+from typing import Optional
+
+import fire  # type: ignore
 import uvicorn
 from mcp.server.fastmcp import FastMCP
 
-from .files import set_workspace_dir
 from .tools.bash import bash
 from .tools.text_editor import text_editor
 from .tools.remote_gpu import (
@@ -11,24 +12,27 @@ from .tools.remote_gpu import (
     create_remote_text_editor,
     remote_download,
 )
+from .files import get_workspace_dir, WorkspaceDirectory
 
 
-server = FastMCP("MLE kit MCP", stateless_http=True)
+def run(host: str = "0.0.0.0", port: int = 5050, workspace: Optional[str] = None) -> None:
+    if workspace:
+        WorkspaceDirectory.set_dir(Path(workspace))
+    workspace_path = get_workspace_dir()
+    workspace_path.mkdir(parents=True, exist_ok=True)
 
-remote_text_editor = create_remote_text_editor(text_editor)
+    server = FastMCP("MLE kit MCP", stateless_http=True)
 
-server.add_tool(bash)
-server.add_tool(text_editor)
-server.add_tool(remote_bash)
-server.add_tool(remote_text_editor)
-server.add_tool(remote_download)
+    remote_text_editor = create_remote_text_editor(text_editor)
 
-http_app = server.streamable_http_app()
+    server.add_tool(bash)
+    server.add_tool(text_editor)
+    server.add_tool(remote_bash)
+    server.add_tool(remote_text_editor)
+    server.add_tool(remote_download)
 
+    http_app = server.streamable_http_app()
 
-def run(host: str = "0.0.0.0", port: int = 5050, workspace: str = "workdir") -> None:
-    workspace_path = Path(workspace)
-    set_workspace_dir(workspace_path)
     uvicorn.run(http_app, host=host, port=port)
 
 

mle_kit_mcp/tools/bash.py

@@ -3,7 +3,7 @@ import atexit
 import signal
 from typing import Optional, Any
 
-from mle_kit_mcp.files import WORKSPACE_DIR_PATH
+from mle_kit_mcp.files import get_workspace_dir
 
 
 _container = None
@@ -63,7 +63,7 @@ def bash(command: str) -> str:
                 tty=True,
                 stdin_open=True,
                 volumes={
-                    WORKSPACE_DIR_PATH: {
+                    get_workspace_dir(): {
                         "bind": DOCKER_WORKSPACE_DIR_PATH,
                         "mode": "rw",
                     }

mle_kit_mcp/tools/remote_gpu.py

@@ -12,7 +12,7 @@ from dataclasses import dataclass
 from dotenv import load_dotenv
 from vastai_sdk import VastAI  # type: ignore
 
-from mle_kit_mcp.files import WORKSPACE_DIR_PATH
+from mle_kit_mcp.files import get_workspace_dir
 
 BASE_IMAGE = "phoenix120/holosophos_mle"
 DEFAULT_GPU_TYPE = "RTX_3090"
@@ -94,7 +94,7 @@ def get_offers(vast_sdk: VastAI, gpu_name: str) -> List[int]:
 
 
 def run_command(
-    instance: InstanceInfo, command: str, timeout: int = 60
+    instance: InstanceInfo, command: str, timeout: int = 60, raise_exc: bool = True
 ) -> subprocess.CompletedProcess[str]:
     cmd = [
         "ssh",
@@ -117,7 +117,7 @@ def run_command(
     ]
     try:
         result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
-        if result.returncode != 0:
+        if result.returncode != 0 and raise_exc:
             raise Exception(
                 f"Error running command: {command}; "
                 f"Output: {result.stdout}; "
@@ -257,9 +257,9 @@ def launch_instance(vast_sdk: VastAI, gpu_name: str) -> Optional[InstanceInfo]:
 
 def send_scripts() -> None:
     assert _instance_info
-    for name in os.listdir(WORKSPACE_DIR_PATH):
+    for name in os.listdir(get_workspace_dir()):
         if name.endswith(".py"):
-            send_rsync(_instance_info, f"{WORKSPACE_DIR_PATH}/{name}", "/root")
+            send_rsync(_instance_info, f"{get_workspace_dir()}/{name}", "/root")
 
 
 def init_all() -> None:
@@ -301,7 +301,7 @@ def remote_bash(command: str, timeout: Optional[int] = 60) -> str:
     init_all()
     assert _instance_info
     assert timeout
-    result = run_command(_instance_info, command, timeout=timeout)
+    result = run_command(_instance_info, command, timeout=timeout, raise_exc=False)
     output = ("STDOUT: " + result.stdout + "\n") if result.stdout else ""
     output += ("STDERR: " + result.stderr) if result.stderr else ""
     return output.replace(VAST_AI_GREETING, "")
@@ -316,7 +316,7 @@ def remote_download(file_path: str) -> str:
     """
     init_all()
     assert _instance_info
-    recieve_rsync(_instance_info, f"/root/{file_path}", f"{WORKSPACE_DIR_PATH}")
+    recieve_rsync(_instance_info, f"/root/{file_path}", f"{get_workspace_dir()}")
     return f"File '{file_path}' downloaded!"
 
 
@@ -335,12 +335,12 @@ def create_remote_text_editor(
         command = args_dict["command"]
 
         if command != "write":
-            recieve_rsync(_instance_info, f"/root/{path}", f"{WORKSPACE_DIR_PATH}")
+            recieve_rsync(_instance_info, f"/root/{path}", f"{get_workspace_dir()}")
 
         result: str = text_editor_func(*args, **kwargs)
 
         if command != "view":
-            send_rsync(_instance_info, f"{WORKSPACE_DIR_PATH}/{path}", "/root")
+            send_rsync(_instance_info, f"{get_workspace_dir()}/{path}", "/root")
 
         return result
 

mle_kit_mcp/tools/text_editor.py

@@ -2,7 +2,7 @@ from collections import defaultdict
 from typing import Dict, List, Optional
 from pathlib import Path
 
-from mle_kit_mcp.files import WORKSPACE_DIR_PATH
+from mle_kit_mcp.files import get_workspace_dir
 from mle_kit_mcp.utils import truncate_content
 
 WRITE_MAX_OUTPUT_LENGTH = 500
@@ -174,7 +174,7 @@ def text_editor(
     ), "Absolute path is not supported, only relative to the work directory"
     valid_commands = ("view", "write", "str_replace", "insert", "undo_edit", "append")
 
-    path_obj = WORKSPACE_DIR_PATH / path
+    path_obj = get_workspace_dir() / path
 
     if command == "view":
         show_lines = show_lines if show_lines is not None else False

{mle_kit_mcp-0.0.2.dist-info → mle_kit_mcp-0.0.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mle-kit-mcp
-Version: 0.0.2
+Version: 0.0.4
 Summary: MCP server that provides different tools for MLE
 Author-email: Ilya Gusev <phoenixilya@gmail.com>
 Project-URL: Homepage, https://github.com/IlyaGusev/mle_kit_mcp

mle_kit_mcp-0.0.4.dist-info/RECORD

@@ -0,0 +1,16 @@
+mle_kit_mcp/__init__.py,sha256=2Ru2I5u4cE7DrkkAsibDUEF1K6sYtqppb9VyFrRoQKI,94
+mle_kit_mcp/__main__.py,sha256=rcmsOtJd3SA82exjrcGBuxuptcoxF8AXI7jNjiVq2BY,59
+mle_kit_mcp/files.py,sha256=ux53kWw7hBAcOmS9qNI4gpQX8XcQPT2LICC--S5-TGI,635
+mle_kit_mcp/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mle_kit_mcp/server.py,sha256=E9cXdKwVBASBzKyrZGHCTvj6BIMN-EbVSQZDFNg0YnE,1062
+mle_kit_mcp/utils.py,sha256=wd7wSyddHRHOYdxmXw8uoAOBxVZOL2_vjNomss07inc,1654
+mle_kit_mcp/tools/__init__.py,sha256=0aLl0gD-JteSvOs2PgVhbv0Wnh6fodFySgQWQvoI1xI,215
+mle_kit_mcp/tools/bash.py,sha256=9IIcely_J_fvNPLdJvHFOL72OG3rVIpNC3MB5S36pJI,2460
+mle_kit_mcp/tools/remote_gpu.py,sha256=Rwy0vuPgBqZDonsJjhapbrBnw1dKViKUJonTC8AiPJE,11429
+mle_kit_mcp/tools/text_editor.py,sha256=s9lPzqQWpZ2277f99TgMyKeOSTrv-PhX4HVtdXVeQ24,9228
+mle_kit_mcp-0.0.4.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mle_kit_mcp-0.0.4.dist-info/METADATA,sha256=-kG6R55hu0Ta3nyNMY87DLWETscN7rMnSelcIISh4SA,1126
+mle_kit_mcp-0.0.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+mle_kit_mcp-0.0.4.dist-info/entry_points.txt,sha256=-iHSUVPN49jkBj1ySpc-P0rVF5-IPHw-KWNayNIiEsk,49
+mle_kit_mcp-0.0.4.dist-info/top_level.txt,sha256=XeBtCq_CnVI0gh0Z_daZOLmGl5XPlkA8RgHaj5s5VQY,12
+mle_kit_mcp-0.0.4.dist-info/RECORD,,

mle_kit_mcp-0.0.2.dist-info/RECORD

@@ -1,16 +0,0 @@
-mle_kit_mcp/__init__.py,sha256=2Ru2I5u4cE7DrkkAsibDUEF1K6sYtqppb9VyFrRoQKI,94
-mle_kit_mcp/__main__.py,sha256=rcmsOtJd3SA82exjrcGBuxuptcoxF8AXI7jNjiVq2BY,59
-mle_kit_mcp/files.py,sha256=loZnpQ7ZLxjX8NgUq-hrvhwbQwmGw1F6gZ3xD3UQI9k,286
-mle_kit_mcp/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mle_kit_mcp/server.py,sha256=l81KAfyZ_G_0Suy9381qV40I_UAZtfT0FKim3_89pSI,886
-mle_kit_mcp/utils.py,sha256=wd7wSyddHRHOYdxmXw8uoAOBxVZOL2_vjNomss07inc,1654
-mle_kit_mcp/tools/__init__.py,sha256=0aLl0gD-JteSvOs2PgVhbv0Wnh6fodFySgQWQvoI1xI,215
-mle_kit_mcp/tools/bash.py,sha256=QcPKQBjD6-LHvNGOAhyWXHuzkM70GATN0Vp85iykSXY,2460
-mle_kit_mcp/tools/remote_gpu.py,sha256=Wh02B-hFUFdRRtZ0J4feLw9P8Hq0AW7xwZ3Q27LcdW4,11370
-mle_kit_mcp/tools/text_editor.py,sha256=zwxyHvf8DdJwuUa7HOWfSbhaSTveaE_WMwsOdDnbq1w,9228
-mle_kit_mcp-0.0.2.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mle_kit_mcp-0.0.2.dist-info/METADATA,sha256=m68NZT-bBF4n1yvKeAfH1arVku2jpkG3juuOs0fQE_4,1126
-mle_kit_mcp-0.0.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-mle_kit_mcp-0.0.2.dist-info/entry_points.txt,sha256=-iHSUVPN49jkBj1ySpc-P0rVF5-IPHw-KWNayNIiEsk,49
-mle_kit_mcp-0.0.2.dist-info/top_level.txt,sha256=XeBtCq_CnVI0gh0Z_daZOLmGl5XPlkA8RgHaj5s5VQY,12
-mle_kit_mcp-0.0.2.dist-info/RECORD,,