Feat: enable B200

Author: S1ro1

src/discord-cluster-manager/modal_runner.py

@@ -19,6 +19,8 @@
 tag = f"{cuda_version}-{flavor}-{operating_sys}"
 
 # Move this to another file later:
+
+
 cuda_image = (
     Image.from_registry(f"nvidia/cuda:{tag}", add_python="3.11")
     .apt_install(
@@ -50,13 +52,20 @@
     .pip_install("requests")
 )
 
+
 cuda_image = cuda_image.add_local_python_source(
     "consts",
     "modal_runner",
     "modal_runner_archs",
     "run_eval",
 )
 
+cuda_image_b200 = (
+    Image.from_registry("nvidia/cuda:12.8.0-devel-ubuntu24.04", add_python="3.11")
+    .pip_install("ninja", "packaging", "requests")
+    .pip_install("torch==2.7.0", extra_index_url="https://download.pytorch.org/whl/cu128")
+)
+
 
 class TimeoutException(Exception):
     pass

src/discord-cluster-manager/modal_runner_archs.py

@@ -1,42 +1,14 @@
 # This file contains wrapper functions for running
 # Modal apps on specific devices. We will fix this later.
-from modal_runner import app, cuda_image, modal_run_config
-from modal_utils import deserialize_full_result
-from run_eval import FullResult, SystemInfo
+from modal_runner import app, cuda_image, cuda_image_b200, modal_run_config
 
 gpus = ["T4", "L4", "A100-80GB", "H100!", "B200", "H200"]
 for gpu in gpus:
     gpu_slug = gpu.lower().split("-")[0].strip("!")
     app.function(gpu=gpu, image=cuda_image, name=f"run_cuda_script_{gpu_slug}", serialized=True)(
         modal_run_config
     )
-    app.function(gpu=gpu, image=cuda_image, name=f"run_pytorch_script_{gpu_slug}", serialized=True)(
+    img = cuda_image if gpu != "B200" else cuda_image_b200
+    app.function(gpu=gpu, image=img, name=f"run_pytorch_script_{gpu_slug}", serialized=True)(
         modal_run_config
     )
-
-
-@app.function(image=cuda_image, max_containers=1, timeout=600)
-def run_pytorch_script_b200(config: dict, timeout: int = 300):
-    """Send a config and timeout to the server and return the response."""
-    import requests
-
-    ip_addr = "34.59.196.5"
-    port = "33001"
-
-    payload = {"config": config, "timeout": timeout}
-
-    try:
-        response = requests.post(f"http://{ip_addr}:{port}", json=payload, timeout=timeout + 5)
-        response.raise_for_status()
-        print("ORIGINAL", response.json())
-
-        print("DESERIALIZED", deserialize_full_result(response.json()))
-        return deserialize_full_result(response.json())
-    except requests.RequestException as e:
-        return FullResult(success=False, error=str(e), runs={}, system=SystemInfo())
-
-
-@app.local_entrypoint()
-def test_b200(timeout: int = 300):
-    config = {}
-    run_pytorch_script_b200.remote(config, timeout)