Update CI, fix GPU index error in zeusd

.github/workflows/deploy_homepage.yaml

@@ -17,7 +17,7 @@ env:
 jobs:
   deploy:
     runs-on: ubuntu-latest
-    if: github.event.repository.fork == false
+    if: github.repository_owner = 'ml-energy'
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4

.github/workflows/publish_crates_io.yaml

@@ -0,0 +1,24 @@
+name: Release
+
+on:
+  push:
+    tags:
+      - zeusd-v*
+
+jobs:
+  cargo-publish:
+    if: github.repository_owner == 'ml-energy'
+    runs-on: ubuntu-latest
+    env:
+      CARGO_TERM_COLOR: always
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          sparse-checkout: zeusd
+      - name: Publish to crates.io
+        uses: katyo/publish-crates@v2
+        with:
+          path: zeusd
+          registry-token: ${{ secrets.CRATES_IO_TOKEN }}
+          check-repo: false

.github/workflows/publish_pypi.yaml

@@ -3,12 +3,12 @@ name: Publish Python package to PyPI
 on:
   push:
     tags:
-      - v*
+      - zeus-v*
 
 jobs:
   publish:
     runs-on: ubuntu-latest
-    if: github.event.repository.fork == false
+    if: github.repository_owner = 'ml-energy'
     steps:
       - name: Checkout repository
         uses: actions/checkout@v3

.github/workflows/push_docker.yaml

@@ -5,7 +5,7 @@ on:
     branches:
       - master
     tags:
-      - v*
+      - zeus-v*
     paths:
       - '.github/workflows/push_docker.yaml'
       - 'capriccio/**'
@@ -21,6 +21,7 @@ on:
 
 jobs:
   build_and_push:
+    if: github.repository_owner == 'ml-energy'
     runs-on: ubuntu-latest
     steps:
       - name: Remove unnecessary files

.github/workflows/fmt_lint_test.yaml → .github/workflows/zeus_fmt_lint_test.yaml

@@ -14,7 +14,7 @@ on:
 
 # Jobs initiated by previous pushes get cancelled by a new push.
 concurrency:
-  group: ${{ github.ref }}-lint-and-test
+  group: ${{ github.ref }}-zeus-lint-and-test
   cancel-in-progress: true
 
 jobs:

zeusd/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
-name = "zeus"
-version = "0.1.0"
+name = "zeusd"
+version = "0.1.1"
 authors = ["Jae-Won Chung <[email protected]>"]
 description = "Zeus daemon"
 license = "Apache-2.0"

zeusd/README.md

@@ -1,3 +1,7 @@
+# Now maintained at [`zeus`](https://crates.io/crates/zeus)
+
+Please contact @jaywonchung if you'd like to take ownership of this handle.
+
 # Zeus daemon (`zeusd`)
 
 `zeusd` is a daemon designed to run with admin privileges and expose API endpoints that wrap privileged NVML methods.
@@ -18,7 +22,7 @@ To make this as low latency as possible, `zeusd` was written with Rust.
 To install `zeusd`:
 
 ```sh
-cargo install zeus
+cargo install zeusd
 ```
 
 With the following, `zeusd` will listen to a unix domain socket at `/var/run/zeusd.sock`, which is writable to anyone (since file permission is 666).

zeusd/src/devices/gpu/macos.rs

@@ -13,7 +13,7 @@ impl NvmlGpu {
 
 impl GpuManager for NvmlGpu {
     fn device_count() -> Result<u32, ZeusdError> {
-        Ok(1)
+        Ok(0)
     }
 
     fn set_persistent_mode(&mut self, _enabled: bool) -> Result<(), ZeusdError> {

zeusd/src/devices/gpu/mod.rs

@@ -98,6 +98,9 @@ impl GpuManagementTasks {
         command: GpuCommand,
         request_start_time: Instant,
     ) -> Result<(), ZeusdError> {
+        if gpu_id >= self.senders.len() {
+            return Err(ZeusdError::GpuNotFoundError(gpu_id));
+        }
         if gpu_id >= self.senders.len() {
             return Err(ZeusdError::GpuNotFoundError(gpu_id));
         }
@@ -114,6 +117,9 @@ impl GpuManagementTasks {
         command: GpuCommand,
         request_start_time: Instant,
     ) -> Result<(), ZeusdError> {
+        if gpu_id >= self.senders.len() {
+            return Err(ZeusdError::GpuNotFoundError(gpu_id));
+        }
         let (tx, mut rx) = tokio::sync::mpsc::channel(1);
         self.senders[gpu_id]
             .send((command, Some(tx), request_start_time, Span::current()))

zeusd/src/main.rs

@@ -41,7 +41,7 @@ async fn main() -> anyhow::Result<()> {
         }
         ConnectionMode::TCP => {
             let listener = TcpListener::bind(&config.tcp_bind_address)?;
-            tracing::info!("Listening on {}", &config.tcp_bind_address);
+            tracing::info!("Listening on {}", &listener.local_addr()?);
 
             start_server_tcp(listener, device_tasks, num_workers)?.await?;
         }

zeusd/tests/gpu.rs

@@ -104,6 +104,20 @@ async fn test_set_persistent_mode_invalid() {
         .await
         .expect("Failed to read response")
         .contains("invalid type"));
+
+    let url = SetPersistentMode::build_url(&app, 5); // Invalid GPU ID
+    let resp = client
+        .post(url)
+        .json(&serde_json::json!(
+            {
+                "enabled": true,
+                "block": true
+            }
+        ))
+        .send()
+        .await
+        .expect("Failed to send request");
+    assert_eq!(resp.status(), 400);
 }
 
 #[tokio::test]
@@ -279,6 +293,20 @@ async fn test_set_power_limit_invalid() {
         .await
         .expect("Failed to read response")
         .contains("missing field"));
+
+    let url = SetPowerLimit::build_url(&app, 5); // Invalid GPU ID
+    let resp = client
+        .post(url)
+        .json(&serde_json::json!(
+            {
+                "power_limit_mw": 100_000,
+                "block": true
+            }
+        ))
+        .send()
+        .await
+        .expect("Failed to send request");
+    assert_eq!(resp.status(), 400);
 }
 
 #[tokio::test]
@@ -466,6 +494,19 @@ async fn test_gpu_locked_clocks_invalid() {
         .await
         .expect("Failed to read response")
         .contains("missing field"));
+
+    let url = ResetGpuLockedClocks::build_url(&app, 5); // Invalid GPU ID
+    let resp = client
+        .post(url)
+        .json(&serde_json::json!(
+            {
+                "block": true
+            }
+        ))
+        .send()
+        .await
+        .expect("Failed to send request");
+    assert_eq!(resp.status(), 400);
 }
 
 #[tokio::test]
@@ -655,6 +696,19 @@ async fn test_mem_locked_clocks_invalid() {
         .await
         .expect("Failed to read response")
         .contains("missing field"));
+
+    let url = ResetMemLockedClocks::build_url(&app, 5); // Invalid GPU ID
+    let resp = client
+        .post(url)
+        .json(&serde_json::json!(
+            {
+                "block": true
+            }
+        ))
+        .send()
+        .await
+        .expect("Failed to send request");
+    assert_eq!(resp.status(), 400);
 }
 
 #[tokio::test]