Iris: Use machine-readable resource quantities in ResourceSpec proto (#2369)

This fixes some layering violations where the controller and worker code
was parsing human-readable strings ("8g") at runtime. By the time
requests reach the controller, they should contain machine-readable
values.

Also removes some typical "claudey" tests for attributes etc.

Author: rjpower

lib/iris/README.md

@@ -66,8 +66,7 @@ src/iris/
 
 ```python
 from iris.client import IrisClient
-from iris.cluster.types import Entrypoint
-from iris.rpc.cluster_pb2 import ResourceSpec
+from iris.cluster.types import Entrypoint, ResourceSpec
 
 def my_task():
     print("Hello from iris!")

lib/iris/examples/demo.ipynb

@@ -50,7 +50,7 @@
     _LocalImageProvider,
 )
 from iris.cluster.controller.controller import Controller, ControllerConfig, DefaultWorkerStubFactory
-from iris.cluster.types import Entrypoint, JobId
+from iris.cluster.types import EnvironmentSpec, Entrypoint, JobId, ResourceSpec
 from iris.cluster.worker.builder import ImageCache
 from iris.cluster.worker.bundle_cache import BundleCache
 from iris.cluster.worker.docker import DockerRuntime
@@ -244,8 +244,8 @@ def submit(
     ) -> str:
         """Submit a job to the cluster."""
         entrypoint = Entrypoint.from_callable(fn, *args, **kwargs)
-        environment = cluster_pb2.EnvironmentConfig(workspace="/app", env_vars={})
-        resources = cluster_pb2.ResourceSpec(cpu=cpu, memory=memory)
+        environment = EnvironmentSpec(workspace="/app")
+        resources = ResourceSpec(cpu=cpu, memory=memory)
         return self.client.submit(
             entrypoint=entrypoint,
             name=name or fn.__name__,

lib/iris/examples/demo_cluster.py

@@ -47,7 +47,7 @@
     RemoteClusterClient,
     get_job_info,
 )
-from iris.cluster.types import Entrypoint, is_job_finished, JobId, Namespace
+from iris.cluster.types import EnvironmentSpec, Entrypoint, is_job_finished, JobId, Namespace, ResourceSpec
 from iris.rpc import cluster_pb2
 from iris.time_utils import ExponentialBackoff
 
@@ -500,8 +500,8 @@ def submit(
         self,
         entrypoint: Entrypoint,
         name: str,
-        resources: cluster_pb2.ResourceSpec,
-        environment: cluster_pb2.EnvironmentConfig | None = None,
+        resources: ResourceSpec,
+        environment: EnvironmentSpec | None = None,
         ports: list[str] | None = None,
         scheduling_timeout_seconds: int = 0,
     ) -> JobId:
@@ -534,11 +534,15 @@ def submit(
         else:
             job_id = JobId(name)
 
+        # Convert to wire format
+        resources_proto = resources.to_proto()
+        environment_proto = environment.to_proto() if environment else None
+
         self._cluster.submit_job(
             job_id=job_id,
             entrypoint=entrypoint,
-            resources=resources,
-            environment=environment,
+            resources=resources_proto,
+            environment=environment_proto,
             ports=ports,
             scheduling_timeout_seconds=scheduling_timeout_seconds,
         )

lib/iris/src/iris/client/client.py

@@ -22,12 +22,13 @@
 Example:
     from pathlib import Path
     from iris.client import IrisClient, WorkerPool, WorkerPoolConfig
+    from iris.cluster.types import ResourceSpec
 
     client = IrisClient.remote("http://controller:8080", workspace=Path("./my-project"))
 
     config = WorkerPoolConfig(
         num_workers=3,
-        resources=cluster_pb2.ResourceSpec(cpu=1, memory="512m"),
+        resources=ResourceSpec(cpu=1, memory="512m"),
     )
 
     with WorkerPool(client, config) as pool:
@@ -55,8 +56,7 @@
 from iris.actor.client import ActorClient
 from iris.actor.resolver import Resolver
 from iris.client.client import IrisClient, iris_ctx
-from iris.cluster.types import Entrypoint, JobId
-from iris.rpc import cluster_pb2
+from iris.cluster.types import EnvironmentSpec, Entrypoint, JobId, ResourceSpec
 from iris.time_utils import ExponentialBackoff
 
 logger = logging.getLogger(__name__)
@@ -329,8 +329,8 @@ class WorkerPoolConfig:
     """
 
     num_workers: int
-    resources: cluster_pb2.ResourceSpec
-    environment: cluster_pb2.EnvironmentConfig | None = None
+    resources: ResourceSpec
+    environment: EnvironmentSpec | None = None
     name_prefix: str = "worker"
     max_retries: int = 0
 

lib/iris/src/iris/client/worker_pool.py

@@ -63,17 +63,16 @@ def __init__(self, cpu: int = 1000, memory_gb: int = 1000):
         self._memory_gb = memory_gb
 
     def probe(self) -> cluster_pb2.WorkerMetadata:
+        device = cluster_pb2.DeviceConfig()
+        device.cpu.CopyFrom(cluster_pb2.CpuDevice(variant="cpu"))
+
         return cluster_pb2.WorkerMetadata(
             hostname="local",
             ip_address="127.0.0.1",
             cpu_count=self._cpu,
             memory_bytes=self._memory_gb * 1024**3,
-        )
-
-    def build_resource_spec(self, metadata: cluster_pb2.WorkerMetadata) -> cluster_pb2.ResourceSpec:
-        return cluster_pb2.ResourceSpec(
-            cpu=metadata.cpu_count,
-            memory=f"{metadata.memory_bytes // (1024**3)}g",
+            disk_bytes=100 * 1024**3,  # Default 100GB for local
+            device=device,
         )
 
 
@@ -368,7 +367,7 @@ def submit_job(
         self,
         job_id: str,
         entrypoint: Entrypoint,
-        resources: cluster_pb2.ResourceSpec,
+        resources: cluster_pb2.ResourceSpecProto,
         environment: cluster_pb2.EnvironmentConfig | None = None,
         ports: list[str] | None = None,
         scheduling_timeout_seconds: int = 0,

lib/iris/src/iris/cluster/client/local_client.py

@@ -67,7 +67,7 @@ def submit_job(
         self,
         job_id: str,
         entrypoint: Entrypoint,
-        resources: cluster_pb2.ResourceSpec,
+        resources: cluster_pb2.ResourceSpecProto,
         environment: cluster_pb2.EnvironmentConfig | None = None,
         ports: list[str] | None = None,
         scheduling_timeout_seconds: int = 0,

lib/iris/src/iris/cluster/client/remote_client.py

@@ -271,9 +271,9 @@ def _dispatch_job(
                     env_vars=dict(job.request.environment.env_vars),
                 ),
                 bundle_gcs_path=job.request.bundle_gcs_path,
-                resources=cluster_pb2.ResourceSpec(
+                resources=cluster_pb2.ResourceSpecProto(
                     cpu=job.request.resources.cpu,
-                    memory=job.request.resources.memory,
+                    memory_bytes=job.request.resources.memory_bytes,
                 ),
                 ports=list(job.request.ports),
             )

lib/iris/src/iris/cluster/controller/controller.py

@@ -222,6 +222,14 @@ def _job_state_name(state: int) -> str:
       return div.innerHTML;
     }
 
+    function formatBytes(bytes) {
+      if (bytes === 0) return '0 B';
+      const k = 1024;
+      const sizes = ['B', 'KB', 'MB', 'GB', 'TB'];
+      const i = Math.floor(Math.log(bytes) / Math.log(k));
+      return parseFloat((bytes / Math.pow(k, i)).toFixed(1)) + ' ' + sizes[i];
+    }
+
     async function refresh() {
       try {
         const [stats, actions, workers, jobs, endpoints] = await Promise.all([
@@ -262,7 +270,8 @@ def _job_state_name(state: int) -> str:
             ? `<a href="http://${escapeHtml(w.address)}/" class="worker-link" target="_blank">${wid}</a>`
             : wid;
           const cpu = w.resources ? w.resources.cpu : '-';
-          const memory = w.resources ? (w.resources.memory || '-') : '-';
+          const memBytes = w.resources ? (w.resources.memory_bytes || 0) : 0;
+          const memory = memBytes ? formatBytes(memBytes) : '-';
           return `<tr>
             <td>${workerLink}</td>
             <td class="${healthClass}">${w.healthy ? 'Yes' : 'No'}</td>
@@ -280,8 +289,9 @@ def _job_state_name(state: int) -> str:
         const jobsHtml = jobs.map(j => {
           const jid = escapeHtml(j.job_id);
           const jobLink = `<a href="/job/${jid}" class="job-link">${jid.slice(0,8)}...</a>`;
+          const jobMemBytes = j.resources ? (j.resources.memory_bytes || 0) : 0;
           const resources = j.resources
-            ? `${j.resources.cpu} CPU, ${j.resources.memory || '-'}` : '-';
+            ? `${j.resources.cpu} CPU, ${jobMemBytes ? formatBytes(jobMemBytes) : '-'}` : '-';
           return `<tr>
             <td>${jobLink}</td>
             <td>${escapeHtml(j.name)}</td>
@@ -778,8 +788,8 @@ def _api_workers(self, _request: Request) -> JSONResponse:
                     "consecutive_failures": w.consecutive_failures,
                     "last_heartbeat_ms": w.last_heartbeat_ms,
                     "resources": {
-                        "cpu": w.resources.cpu if w.resources else 0,
-                        "memory": w.resources.memory if w.resources else "",
+                        "cpu": w.metadata.cpu_count,
+                        "memory_bytes": w.metadata.memory_bytes,
                     },
                 }
                 for w in workers
@@ -805,7 +815,7 @@ def _api_jobs(self, _request: Request) -> JSONResponse:
                     "preemption_count": j.preemption_count,
                     "resources": {
                         "cpu": j.request.resources.cpu if j.request.resources else 0,
-                        "memory": j.request.resources.memory if j.request.resources else "",
+                        "memory_bytes": j.request.resources.memory_bytes if j.request.resources else 0,
                     },
                 }
                 for j in jobs