Bug fix: cpu metrics missing functionality restored (#12)

* bug fix: missing functions

* mps device support

* bump version

* return formatting to ruff standard

* silence ruff checks in locations that could impede functionality if we do not cleanly fix.

* add pytest to dev dependencies

* add other requirements files for cuda versions

* correct typo in Readme

---------

Co-authored-by: nbovee <12849851+nbovee@users.noreply.github.com>

Author: ali-izhar

README.md

@@ -44,7 +44,7 @@
    ```bash
    git clone https://github.com/nbovee/tracr.git && cd tracr
    python3 -m venv venv && source venv/bin/activate
-   pip install -r requirements.txt
+   pip install -r requirements.txt # alternatively, use the requirements-cu###.txt file for your cuda version.
    ```
 
 2. **Configure devices** by copying and editing the template
@@ -327,7 +327,10 @@ Select optimal split points based on:
 <summary>Unit Testing</summary>
 - If issues present themselves, the provided unit tests may have some insight to the error. Please run the following, and refine to individual files for further details:
   
-  ```python -m unittest discovery -s ./tests```
+  ```python -m unittest discover -s ./tests```
+  or if using uv, the command will be:
+
+  ```uv run -m unittest discover -s ./tests```
 </details>
 
 <summary>Connection Issues</summary>

config/modelsplit_template.yaml

@@ -13,7 +13,7 @@
 # DEFAULT CONFIGURATIONS
 # ================================================================
 default:
-  device: "cuda"                      # [OPTIONAL] Computing device: 'cuda' (GPU) or 'cpu'. Default: 'cuda' if available, else 'cpu'
+  device: "cuda"                      # [OPTIONAL] Computing device: 'cuda' (NVIDIA GPU), 'mps' (Apple Silicon GPU), or 'cpu'. Default: 'cuda' if available, 'mps' on Apple Silicon if available, else 'cpu'
   save_layer_images: false            # [OPTIONAL] Save intermediate layer images. Default: false
   collect_metrics: false              # [OPTIONAL] Collect detailed metrics per layer (time-consuming for cpu systems). Default: false
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tracr"
-version = "0.4.0"
+version = "0.4.1"
 description = "An experimental framework for computational offloading and distributed neural network inference through split computing"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -61,6 +61,11 @@ conflicts = [
 [tool.ruff.lint]
 ignore = ["F841"]
 
+[dependency-groups]
+dev = [
+    "pytest>=8.3.5",
+]
+
 [tool.uv.sources]
 torch = [
   { index = "pytorch-cpu", extra = "cpu", marker = "platform_system != 'Darwin'" },

requirements-cu118.txt

@@ -0,0 +1,30 @@
+# This file was autogenerated by uv via the following command:
+#    uv pip compile pyproject.toml -o requirements-cu118.txt --no-deps --extra cu118 --extra full
+blosc2==3.2.0
+    # via tracr (pyproject.toml)
+loguru==0.7.3
+    # via tracr (pyproject.toml)
+pandas==2.2.3
+    # via tracr (pyproject.toml)
+paramiko==3.5.1
+    # via tracr (pyproject.toml)
+pyyaml==6.0.2
+    # via tracr (pyproject.toml)
+rich==13.9.4
+    # via tracr (pyproject.toml)
+rpyc==6.0.1
+    # via tracr (pyproject.toml)
+tomli==2.2.1
+    # via tracr (pyproject.toml)
+torch==2.6.0+cu118
+    # via tracr (pyproject.toml)
+torchaudio==2.6.0+cu118
+    # via tracr (pyproject.toml)
+torchinfo==1.8.0
+    # via tracr (pyproject.toml)
+torchvision==0.21.0+cu118
+    # via tracr (pyproject.toml)
+tqdm==4.67.1
+    # via tracr (pyproject.toml)
+ultralytics==8.3.93
+    # via tracr (pyproject.toml)

requirements-cu121.txt

@@ -0,0 +1,30 @@
+# This file was autogenerated by uv via the following command:
+#    uv pip compile pyproject.toml -o requirements-cu121.txt --no-deps --extra cu121 --extra full
+blosc2==3.2.0
+    # via tracr (pyproject.toml)
+loguru==0.7.3
+    # via tracr (pyproject.toml)
+pandas==2.2.3
+    # via tracr (pyproject.toml)
+paramiko==3.5.1
+    # via tracr (pyproject.toml)
+pyyaml==6.0.2
+    # via tracr (pyproject.toml)
+rich==13.9.4
+    # via tracr (pyproject.toml)
+rpyc==6.0.1
+    # via tracr (pyproject.toml)
+tomli==2.2.1
+    # via tracr (pyproject.toml)
+torch==2.5.1+cu121
+    # via tracr (pyproject.toml)
+torchaudio==2.5.1+cu121
+    # via tracr (pyproject.toml)
+torchinfo==1.8.0
+    # via tracr (pyproject.toml)
+torchvision==0.20.1+cu121
+    # via tracr (pyproject.toml)
+tqdm==4.67.1
+    # via tracr (pyproject.toml)
+ultralytics==8.3.93
+    # via tracr (pyproject.toml)

requirements-cu124.txt

@@ -0,0 +1,30 @@
+# This file was autogenerated by uv via the following command:
+#    uv pip compile pyproject.toml -o requirements-cu124.txt --no-deps --extra cu124 --extra full
+blosc2==3.2.0
+    # via tracr (pyproject.toml)
+loguru==0.7.3
+    # via tracr (pyproject.toml)
+pandas==2.2.3
+    # via tracr (pyproject.toml)
+paramiko==3.5.1
+    # via tracr (pyproject.toml)
+pyyaml==6.0.2
+    # via tracr (pyproject.toml)
+rich==13.9.4
+    # via tracr (pyproject.toml)
+rpyc==6.0.1
+    # via tracr (pyproject.toml)
+tomli==2.2.1
+    # via tracr (pyproject.toml)
+torch==2.6.0+cu124
+    # via tracr (pyproject.toml)
+torchaudio==2.6.0+cu124
+    # via tracr (pyproject.toml)
+torchinfo==1.8.0
+    # via tracr (pyproject.toml)
+torchvision==0.21.0+cu124
+    # via tracr (pyproject.toml)
+tqdm==4.67.1
+    # via tracr (pyproject.toml)
+ultralytics==8.3.93
+    # via tracr (pyproject.toml)

server.py

@@ -25,17 +25,17 @@
 if str(project_root) not in sys.path:
     sys.path.append(str(project_root))
 
-from src.api import (
+from src.api import ( # noqa: E402
     DataCompression,
     DeviceManager,
     ExperimentManager,
     DeviceType,
     start_logging_server,
     shutdown_logging_server,
-    DataCompression,
+    DataCompression, # noqa: F811
     read_yaml_file,
 )
-from src.api.network.protocols import (
+from src.api.network.protocols import ( # noqa: E402
     LENGTH_PREFIX_SIZE,
     ACK_MESSAGE,
     SERVER_COMPRESSION_SETTINGS,
@@ -61,11 +61,43 @@ def get_device(requested_device: str = "cuda") -> str:
         logger.info("CPU device explicitly requested")
         return "cpu"
 
-    if requested_device in ("cuda", "gpu", "mps") and torch.cuda.is_available():
+    if requested_device == "cuda" and torch.cuda.is_available():
         logger.info("CUDA is available and will be used")
         return "cuda"
 
-    logger.warning("CUDA requested but not available, falling back to CPU")
+    # Check for MPS (Apple Silicon GPUs)
+    if (
+        requested_device == "mps"
+        and hasattr(torch.backends, "mps")
+        and torch.backends.mps.is_available()
+    ):
+        logger.info("MPS (Apple Silicon GPU) is available and will be used")
+        return "mps"
+
+    # If we're here, requested GPU is not available - try alternatives
+    if requested_device in ("cuda", "gpu", "mps"):
+        # If any GPU was requested, try all available options in priority order
+        if torch.cuda.is_available():
+            logger.info(
+                f"{requested_device.upper()} requested but not available, using CUDA instead"
+            )
+            return "cuda"
+
+        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            logger.info(
+                f"{requested_device.upper()} requested but not available, using MPS (Apple Silicon GPU) instead"
+            )
+            return "mps"
+
+        logger.warning(
+            f"{requested_device.upper()} requested but no GPU available, falling back to CPU"
+        )
+        return "cpu"
+
+    # For any other requested device, fall back to CPU
+    logger.warning(
+        f"Requested device '{requested_device}' not recognized, falling back to CPU"
+    )
     return "cpu"
 
 
@@ -353,7 +385,7 @@ def _receive_config(self, conn: socket.socket) -> dict:
             # Deserialize using pickle
             try:
                 config = pickle.loads(config_data)
-                logger.debug(f"Successfully received and parsed configuration")
+                logger.debug("Successfully received and parsed configuration")
                 return config
             except Exception as e:
                 logger.error(f"Failed to deserialize config: {e}")

src/api/__init__.py

@@ -1,13 +1,13 @@
 """API module"""
 
-from .core import *
-from .devices import *
-from .experiments import *
-from .inference import *
-from .network import *
-from .utils import *
+from .core import * # noqa: F403
+from .devices import * # noqa: F403
+from .experiments import * # noqa: F403
+from .inference import * # noqa: F403
+from .network import * # noqa: F403
+from .utils import * # noqa: F403
 
-__all__ = [
+__all__ = [ # noqa: F405
     "core",
     "devices",
     "experiments",

src/api/devices/manager.py

@@ -16,7 +16,7 @@
 )
 from .discovery import LAN
 from ..network.ssh import SSHKeyHandler, SSHConfig, create_ssh_client
-from ..network.protocols import SSH_PORT, SSH_CONNECTIVITY_TIMEOUT, DEFAULT_PORT
+from ..network.protocols import SSH_PORT, SSH_CONNECTIVITY_TIMEOUT, DEFAULT_PORT # noqa: F401
 from ..utils.utils import get_repo_root
 
 logger = logging.getLogger("split_computing_logger")

src/api/experiments/base.py

@@ -1,7 +1,7 @@
 """Core experiment infrastructure for split computing"""
 
 import logging
-import sys
+import sys # noqa: F401
 import time
 from dataclasses import dataclass, field
 from pathlib import Path