mle-kit-mcp 0.1.2__py3-none-any.whl → 0.2.0__py3-none-any.whl

mle_kit_mcp/tools/llm_proxy.py

@@ -116,7 +116,7 @@ def llm_proxy_remote(port: Optional[int] = None) -> str:
     chosen_port = port or random.randint(5000, 6000)
     token = secrets.token_urlsafe(24)
     dependencies_cmd = f"python3 -m pip install -q --no-input {DEPENDENCIES}"
-    _remote_run_command(instance, dependencies_cmd, timeout=300, raise_exc=True)
+    _remote_run_command(instance, dependencies_cmd, timeout=300)
 
     launch_cmd = (
         f"OPENROUTER_API_KEY='{api_key}' ACCESS_TOKEN='{token}' "
@@ -125,14 +125,12 @@ def llm_proxy_remote(port: Optional[int] = None) -> str:
         f"> openrouter_proxy.log 2>&1 "
         f"& echo $! > openrouter_proxy.pid"
     )
-    _remote_run_command(instance, launch_cmd, timeout=60, raise_exc=True)
+    _remote_run_command(instance, launch_cmd, timeout=60)
 
     health_cmd = f'import httpx; print(httpx.get("http://127.0.0.1:{chosen_port}/health").json())'
     start_time = time.time()
     while time.time() - start_time < START_TIMEOUT:
-        result = _remote_run_command(
-            instance, f"python -c '{health_cmd}'", timeout=10, raise_exc=False
-        )
+        result = _remote_run_command(instance, f"python -c '{health_cmd}'", timeout=10)
         if result.returncode == 0 and "ok" in result.stdout.strip():
             break
         time.sleep(1)

mle_kit_mcp/tools/remote_gpu.py

@@ -115,7 +115,7 @@ def get_offers(vast_sdk: VastAI, gpu_name: str) -> List[int]:
 
 
 def run_command(
-    instance: InstanceInfo, command: str, timeout: int = 60, raise_exc: bool = True
+    instance: InstanceInfo, command: str, timeout: int = 60
 ) -> subprocess.CompletedProcess[str]:
     cmd = [
         "ssh",
@@ -138,16 +138,31 @@ def run_command(
     ]
     try:
         result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
-        if result.returncode != 0 and raise_exc:
-            raise Exception(
-                f"Error running command: {command}; "
-                f"Output: {result.stdout}; "
-                f"Error: {result.stderr}"
-            )
-    except subprocess.TimeoutExpired:
-        raise Exception(
+        if result.stdout:
+            result.stdout = result.stdout.replace(VAST_AI_GREETING, "").strip()
+        if result.stderr:
+            result.stderr = result.stderr.replace(VAST_AI_GREETING, "").strip()
+    except subprocess.TimeoutExpired as e:
+        output = None
+        if e.stdout:
+            output_message = e.stdout.decode("utf-8").strip()
+            output_message = output_message.replace(VAST_AI_GREETING, "").strip()
+            if output_message:
+                output = output_message
+        error = (
             f"Command timed out after {timeout} seconds: {command};\n"
-            f"You can increase the timeout by changing the parameter of the tool call."
+            f"You can increase the timeout by changing the parameter of the tool call.\n"
+        )
+        if e.stderr:
+            error_message = e.stderr.decode("utf-8").strip()
+            error_message = error_message.replace(VAST_AI_GREETING, "").strip()
+            if error_message:
+                error += f"Original stderr: {error_message}"
+        return subprocess.CompletedProcess(
+            args=cmd,
+            returncode=-9,
+            stdout=output,
+            stderr=error,
         )
     return result
 
@@ -255,10 +270,9 @@ def launch_instance(vast_sdk: VastAI, gpu_name: str) -> Optional[InstanceInfo]:
         max_attempts = 10
         is_okay = False
         for attempt in range(max_attempts):
-            try:
-                result = run_command(info, "echo 'SSH connection successful'")
-            except Exception as e:
-                print(f"Waiting for SSH... {e}\n(Attempt {attempt+1}/{max_attempts})")
+            result = run_command(info, "echo 'SSH connection successful'")
+            if result.returncode != 0:
+                print(f"Waiting for SSH... {result.stderr}\n(Attempt {attempt+1}/{max_attempts})")
                 time.sleep(30)
                 continue
             if "SSH connection successful" in result.stdout:
@@ -305,10 +319,10 @@ def remote_bash(command: str, timeout: int = 60) -> str:
     instance = get_instance()
     assert instance
     assert timeout
-    result = run_command(instance, 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, "")
+    result = run_command(instance, command, timeout=timeout)
+    output = ("Command stdout: " + result.stdout + "\n") if result.stdout else ""
+    output += ("Command stderr: " + result.stderr) if result.stderr else ""
+    return output
 
 
 def remote_download(file_path: str) -> str:

{mle_kit_mcp-0.1.2.dist-info → mle_kit_mcp-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mle-kit-mcp
-Version: 0.1.2
+Version: 0.2.0
 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.1.2.dist-info → mle_kit_mcp-0.2.0.dist-info}/RECORD

@@ -7,12 +7,12 @@ mle_kit_mcp/server.py,sha256=W4YJ3m1-NKheJ5QlGggfMQyXPghzKgg5zXqb-TLYH1U,1271
 mle_kit_mcp/utils.py,sha256=iHNcEZZzPD37bEYE18SzJ3WUjLP3Ym-kc91SwcW1vlI,1984
 mle_kit_mcp/tools/__init__.py,sha256=r2fIg2mZ6zaeq0CzEKCEdeUTjV0pcA9NZaaOfBNVTnE,332
 mle_kit_mcp/tools/bash.py,sha256=kunYHc3dyPGOooT-KY9L7eI_N22lBrcDbTlcp_yTTws,2820
-mle_kit_mcp/tools/llm_proxy.py,sha256=eV8QwDBgsyfBzD454GzWsNjH6976E-e_fmO6CYpdLGc,5305
-mle_kit_mcp/tools/remote_gpu.py,sha256=Fv8SJC8bE3Oo8JHOKXoYpISXFMzUMFjEcCQOTyW7V1I,11544
+mle_kit_mcp/tools/llm_proxy.py,sha256=RHW-PgxNnq2rs3sLJCA_CHNGt2ar1DYN2U0OlRjvDZ0,5234
+mle_kit_mcp/tools/remote_gpu.py,sha256=DgfRkMVKn1HYLLboo3nkVIy1qcTlsX4xADQCAHckp_I,12130
 mle_kit_mcp/tools/text_editor.py,sha256=0uGcSjcBjdPlqjYZYuWo29SmAMzBcWgKt4KRumFI3WE,9529
-mle_kit_mcp-0.1.2.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mle_kit_mcp-0.1.2.dist-info/METADATA,sha256=pGh_SbX6SFp_Ku1aOIUmPQjpDFXvBeXFebcyVzQcJVM,1074
-mle_kit_mcp-0.1.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-mle_kit_mcp-0.1.2.dist-info/entry_points.txt,sha256=-iHSUVPN49jkBj1ySpc-P0rVF5-IPHw-KWNayNIiEsk,49
-mle_kit_mcp-0.1.2.dist-info/top_level.txt,sha256=XeBtCq_CnVI0gh0Z_daZOLmGl5XPlkA8RgHaj5s5VQY,12
-mle_kit_mcp-0.1.2.dist-info/RECORD,,
+mle_kit_mcp-0.2.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mle_kit_mcp-0.2.0.dist-info/METADATA,sha256=BbJCb7Y1ep9BIeJk5MpEOEqxSCf3RRd3vBpCqJZDbuM,1074
+mle_kit_mcp-0.2.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+mle_kit_mcp-0.2.0.dist-info/entry_points.txt,sha256=-iHSUVPN49jkBj1ySpc-P0rVF5-IPHw-KWNayNIiEsk,49
+mle_kit_mcp-0.2.0.dist-info/top_level.txt,sha256=XeBtCq_CnVI0gh0Z_daZOLmGl5XPlkA8RgHaj5s5VQY,12
+mle_kit_mcp-0.2.0.dist-info/RECORD,,