diff --git a/.gitignore b/.gitignore index 72fd11b6a..47139bf23 100644 --- a/.gitignore +++ b/.gitignore @@ -35,6 +35,7 @@ sandbox/kafka-data sandbox/zookeeper-data sandbox/zookeeper-logs sandbox/rqd/shots/ +sandbox/pgadmin-data docs/_data/version.yml target/* diff --git a/VERSION.in b/VERSION.in index b48f32260..adc97d8e2 100644 --- a/VERSION.in +++ b/VERSION.in @@ -1 +1 @@ -1.17 +1.18 diff --git a/cuebot/build.gradle b/cuebot/build.gradle index 6e30c1e8c..e7242db19 100644 --- a/cuebot/build.gradle +++ b/cuebot/build.gradle @@ -89,6 +89,7 @@ compileTestJava { options.compilerArgs << "-Xlint:all,-serial" << "-Werror" } + protobuf { protoc { // The protoc compiler diff --git a/cuebot/src/main/java/com/imageworks/spcue/DispatchFrame.java b/cuebot/src/main/java/com/imageworks/spcue/DispatchFrame.java index b73cd75c3..a006a7a79 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/DispatchFrame.java +++ b/cuebot/src/main/java/com/imageworks/spcue/DispatchFrame.java @@ -43,6 +43,7 @@ public class DispatchFrame extends FrameEntity implements FrameInterface { public int minGpus; public int maxGpus; public long minGpuMemory; + public int slotsRequired; // A comma separated list of services public String services; diff --git a/cuebot/src/main/java/com/imageworks/spcue/DispatchHost.java b/cuebot/src/main/java/com/imageworks/spcue/DispatchHost.java index a4bee8e75..43473ac12 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/DispatchHost.java +++ b/cuebot/src/main/java/com/imageworks/spcue/DispatchHost.java @@ -49,6 +49,7 @@ public class DispatchHost extends Entity public long idleGpuMemory; public String tags; private String os; + public int runningSlots; public boolean isNimby; public boolean isLocalDispatch = false; diff --git a/cuebot/src/main/java/com/imageworks/spcue/HostEntity.java b/cuebot/src/main/java/com/imageworks/spcue/HostEntity.java index 548e99cd2..c965513b6 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/HostEntity.java +++ b/cuebot/src/main/java/com/imageworks/spcue/HostEntity.java @@ -38,6 +38,8 @@ public class HostEntity extends Entity implements HostInterface { public int idleGpus; public long gpuMemory; public long idleGpuMemory; + public int concurrentSlotsLimit; + public int runningSlots; public boolean unlockAtBoot; @@ -61,6 +63,8 @@ public HostEntity(Host grpcHost) { this.idleGpus = (int) grpcHost.getIdleGpus(); this.gpuMemory = grpcHost.getGpuMemory(); this.idleGpuMemory = grpcHost.getIdleGpuMemory(); + this.concurrentSlotsLimit = grpcHost.getConcurrentSlotsLimit(); + this.runningSlots = grpcHost.getRunningSlots(); } public String getHostId() { diff --git a/cuebot/src/main/java/com/imageworks/spcue/LayerDetail.java b/cuebot/src/main/java/com/imageworks/spcue/LayerDetail.java index 572139039..da0e830f6 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/LayerDetail.java +++ b/cuebot/src/main/java/com/imageworks/spcue/LayerDetail.java @@ -38,6 +38,7 @@ public class LayerDetail extends LayerEntity implements LayerInterface { public int timeout_llu; public int dispatchOrder; public int totalFrameCount; + public int slotsRequired; public Set tags = new LinkedHashSet(); public Set services = new LinkedHashSet(); diff --git a/cuebot/src/main/java/com/imageworks/spcue/VirtualProc.java b/cuebot/src/main/java/com/imageworks/spcue/VirtualProc.java index 8c92ad016..4f953ee5e 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/VirtualProc.java +++ b/cuebot/src/main/java/com/imageworks/spcue/VirtualProc.java @@ -45,6 +45,8 @@ public class VirtualProc extends FrameEntity implements ProcInterface { public long gpuMemoryUsed; public long gpuMemoryMax; + public int slotsRequired; + public boolean unbooked; public boolean usageRecorded = false; public boolean isLocalDispatch = false; @@ -101,6 +103,7 @@ public static final VirtualProc build(DispatchHost host, DispatchFrame frame, proc.memoryReserved = frame.getMinMemory(); proc.gpusReserved = frame.minGpus; proc.gpuMemoryReserved = frame.minGpuMemory; + proc.slotsRequired = frame.slotsRequired; /* * Frames that are announcing cores less than 100 are not multi-threaded so there is no @@ -237,6 +240,7 @@ public static final VirtualProc build(DispatchHost host, DispatchFrame frame, proc.memoryReserved = frame.getMinMemory(); proc.gpusReserved = frame.minGpus; proc.gpuMemoryReserved = frame.minGpuMemory; + proc.slotsRequired = frame.slotsRequired; int wholeCores = (int) (Math.floor(host.idleCores / 100.0)); if (wholeCores == 0) { diff --git a/cuebot/src/main/java/com/imageworks/spcue/dao/HostDao.java b/cuebot/src/main/java/com/imageworks/spcue/dao/HostDao.java index dfd0397a2..9353c46aa 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/dao/HostDao.java +++ b/cuebot/src/main/java/com/imageworks/spcue/dao/HostDao.java @@ -243,6 +243,22 @@ public interface HostDao { */ void updateThreadMode(HostInterface host, ThreadMode mode); + /** + * Update the host's concurrent procs limit. + * + * @param host HostInterface + * @param limit int (0 for no limit) + */ + void updateConcurrentSlotsLimit(HostInterface host, int limit); + + /** + * Get the host's concurrent slots limit by hostname. + * + * @param hostname String + * @return int the concurrent slots limit + */ + int getHostConcurrentSlotsLimit(String hostname); + /** * Update the specified host's hardware information. * @@ -257,10 +273,11 @@ public interface HostDao { * @param freeGpuMemory long * @param load int * @param os String + * @param runningSlots int */ void updateHostStats(HostInterface host, long totalMemory, long freeMemory, long totalSwap, long freeSwap, long totalMcp, long freeMcp, long totalGpuMemory, long freeGpuMemory, - int load, Timestamp bootTime, String os); + int load, Timestamp bootTime, String os, int runningSlots); /** * Return true if the HardwareState is Up, false if it is anything else. diff --git a/cuebot/src/main/java/com/imageworks/spcue/dao/LayerDao.java b/cuebot/src/main/java/com/imageworks/spcue/dao/LayerDao.java index 847ae231e..8d5988706 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/dao/LayerDao.java +++ b/cuebot/src/main/java/com/imageworks/spcue/dao/LayerDao.java @@ -321,6 +321,14 @@ public interface LayerDao { */ void updateTimeoutLLU(LayerInterface layer, int timeout_llu); + /** + * Updates the slots required for a layer. + * + * @param layer the layer to update + * @param slots the number of slots required (<0 means the host is not slot-based) + */ + void updateLayerSlotsRequired(LayerInterface layer, int slots); + /** * Lowers the minimum memory on a layer if the layer is using less memory and the currnet min * memory is the dispatcher default. diff --git a/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/DispatchQuery.java b/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/DispatchQuery.java index 2daa00335..4f75d6708 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/DispatchQuery.java +++ b/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/DispatchQuery.java @@ -546,6 +546,7 @@ private static final String replaceQueryForFifo(String query) { "int_gpus_min, " + "int_gpus_max, " + "int_gpu_mem_min, " + + "int_slots_required, " + "str_cmd, " + "str_range, " + "int_chunk_size, " + @@ -588,6 +589,7 @@ private static final String replaceQueryForFifo(String query) { "layer.int_gpus_min, " + "layer.int_gpus_max, " + "layer.int_gpu_mem_min, " + + "layer.int_slots_required, " + "layer.str_cmd, " + "layer.str_range, " + "layer.int_chunk_size, " + @@ -676,6 +678,7 @@ private static final String replaceQueryForFifo(String query) { "layer.b_threadable, " + "layer.int_mem_min, " + "layer.int_gpu_mem_min, " + + "layer.int_slots_required, " + "layer.str_cmd, " + "layer.str_range, " + "layer.int_chunk_size, " + @@ -765,6 +768,7 @@ private static final String replaceQueryForFifo(String query) { "layer.int_gpus_min, " + "layer.int_gpus_max, " + "layer.int_gpu_mem_min, " + + "layer.int_slots_required, " + "layer.str_cmd, " + "layer.str_range, " + "layer.int_chunk_size, " + @@ -847,6 +851,7 @@ private static final String replaceQueryForFifo(String query) { "layer.int_gpus_min, " + "layer.int_gpus_max, " + "layer.int_gpu_mem_min, " + + "layer.int_slots_required, " + "layer.str_cmd, " + "layer.str_range, " + "layer.int_chunk_size, " + @@ -932,6 +937,7 @@ private static final String replaceQueryForFifo(String query) { "layer.int_gpus_min, " + "layer.int_gpus_max, " + "layer.int_gpu_mem_min, " + + "layer.int_slots_required, " + "layer.str_cmd, " + "layer.str_range, " + "layer.int_chunk_size, " + @@ -1020,6 +1026,7 @@ private static final String replaceQueryForFifo(String query) { "layer.int_gpus_min, " + "layer.int_gpus_max, " + "layer.int_gpu_mem_min, " + + "layer.int_slots_required, " + "layer.str_cmd, " + "layer.str_range, " + "layer.int_chunk_size, " + @@ -1108,6 +1115,7 @@ private static final String replaceQueryForFifo(String query) { "layer.int_gpus_min, " + "layer.int_gpus_max, " + "layer.int_gpu_mem_min, " + + "layer.int_slots_required, " + "layer.int_cores_max, " + "layer.str_cmd, " + "layer.str_range, " + @@ -1191,6 +1199,7 @@ private static final String replaceQueryForFifo(String query) { "layer.int_gpus_min, " + "layer.int_gpus_max, " + "layer.int_gpu_mem_min, " + + "layer.int_slots_required, " + "layer.str_cmd, " + "layer.str_range, " + "layer.int_chunk_size, " + diff --git a/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/FrameDaoJdbc.java b/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/FrameDaoJdbc.java index 62e14746f..2fa736c7c 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/FrameDaoJdbc.java +++ b/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/FrameDaoJdbc.java @@ -234,6 +234,7 @@ public DispatchFrame mapRow(ResultSet rs, int rowNum) throws SQLException { frame.minGpus = rs.getInt("int_gpus_min"); frame.maxGpus = rs.getInt("int_gpus_max"); frame.minGpuMemory = rs.getLong("int_gpu_mem_min"); + frame.slotsRequired = rs.getInt("int_slots_required"); frame.version = rs.getInt("int_version"); frame.services = rs.getString("str_services"); frame.os = rs.getString("str_os"); @@ -252,8 +253,8 @@ public DispatchFrame mapRow(ResultSet rs, int rowNum) throws SQLException { + "layer.str_type AS layer_type, " + "layer.str_cmd, " + "layer.int_cores_min," + "layer.int_cores_max," + "layer.b_threadable," + "layer.int_mem_min, " + "layer.int_gpus_min," + "layer.int_gpus_max," + "layer.int_gpu_mem_min, " - + "layer.str_range, " + "layer.int_chunk_size, " + "layer.str_services " + "FROM " - + "layer, " + "job, " + "show, " + + "layer.int_slots_required, " + "layer.str_range, " + "layer.int_chunk_size, " + + "layer.str_services " + "FROM " + "layer, " + "job, " + "show, " + "frame LEFT JOIN proc ON (proc.pk_frame = frame.pk_frame) " + "WHERE " + "job.pk_show = show.pk_show " + "AND " + "frame.pk_job = job.pk_job " + "AND " + "frame.pk_layer = layer.pk_layer " + "AND " + "frame.pk_frame = ?"; diff --git a/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/HostDaoJdbc.java b/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/HostDaoJdbc.java index 862e1c459..ae3831311 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/HostDaoJdbc.java +++ b/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/HostDaoJdbc.java @@ -78,6 +78,8 @@ public HostEntity mapRow(ResultSet rs, int rowNum) throws SQLException { host.idleGpus = rs.getInt("int_gpus_idle"); host.gpuMemory = rs.getLong("int_gpu_mem"); host.idleGpuMemory = rs.getLong("int_gpu_mem_idle"); + host.concurrentSlotsLimit = rs.getInt("int_concurrent_slots_limit"); + host.runningSlots = rs.getInt("int_running_slots"); host.dateBooted = rs.getDate("ts_booted"); host.dateCreated = rs.getDate("ts_created"); host.datePinged = rs.getDate("ts_ping"); @@ -131,11 +133,13 @@ public String getFacilityId() { + " host.int_gpus_idle, " + " host.int_gpu_mem, " + " host.int_gpu_mem_idle, " + + " host.int_concurrent_slots_limit, " + " host.ts_created, " + " host.str_name, " + " host_stat.str_state, " + " host_stat.ts_ping, " + " host_stat.ts_booted, " + + " host_stat.int_running_slots, " + " alloc.pk_facility " + "FROM " + " host, " @@ -229,6 +233,7 @@ public DispatchHost mapRow(ResultSet rs, int rowNum) throws SQLException { host.tags = rs.getString("str_tags"); host.setOs(rs.getString("str_os")); host.hardwareState = HardwareState.valueOf(rs.getString("str_state")); + host.runningSlots = rs.getInt("int_running_slots"); return host; } }; @@ -253,6 +258,7 @@ public DispatchHost mapRow(ResultSet rs, int rowNum) throws SQLException { + " host.str_tags, " + " host_stat.str_os, " + " host_stat.str_state, " + + " host_stat.int_running_slots, " + " alloc.pk_facility " + "FROM " + " host " @@ -395,14 +401,15 @@ public CallableStatement createCallableStatement(Connection con) throws SQLExcep + " int_load = ?, " + " ts_booted = ?, " + " ts_ping = current_timestamp, " - + " str_os = ? " + + " str_os = ?, " + + " int_running_slots = ? " + "WHERE " + " pk_host = ?"; @Override public void updateHostStats(HostInterface host, long totalMemory, long freeMemory, long totalSwap, long freeSwap, long totalMcp, long freeMcp, long totalGpuMemory, - long freeGpuMemory, int load, Timestamp bootTime, String os) { + long freeGpuMemory, int load, Timestamp bootTime, String os, int runningSlots) { if (os == null) { os = Dispatcher.OS_DEFAULT; @@ -410,7 +417,7 @@ public void updateHostStats(HostInterface host, long totalMemory, long freeMemor getJdbcTemplate().update(UPDATE_RENDER_HOST, totalMemory, freeMemory, totalSwap, freeSwap, totalMcp, freeMcp, totalGpuMemory, freeGpuMemory, load, bootTime, os, - host.getHostId()); + runningSlots, host.getHostId()); } @Override @@ -562,6 +569,23 @@ public void updateThreadMode(HostInterface host, ThreadMode mode) { mode.getNumber(), host.getHostId()); } + @Override + public void updateConcurrentSlotsLimit(HostInterface host, int limit) { + getJdbcTemplate().update("UPDATE host SET int_concurrent_slots_limit=? WHERE pk_host=?", + limit, host.getHostId()); + } + + @Override + public int getHostConcurrentSlotsLimit(String hostname) { + try { + return getJdbcTemplate().queryForObject( + "SELECT int_concurrent_slots_limit FROM host WHERE str_name = ?", + Integer.class, hostname); + } catch (EmptyResultDataAccessException e) { + return 0; + } + } + @Override public void updateHostOs(HostInterface host, String os) { getJdbcTemplate().update("UPDATE host_stat SET str_os=? WHERE pk_host=?", os, @@ -631,7 +655,7 @@ public boolean isNimbyHost(HostInterface h) { /** * Checks if the passed in name looks like a fully qualified domain name. If so, returns the * hostname without the domain. Otherwise returns the passed in name unchanged. - * + * * @param fqdn - String * @return String - hostname */ diff --git a/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/LayerDaoJdbc.java b/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/LayerDaoJdbc.java index 910999a48..2991a3640 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/LayerDaoJdbc.java +++ b/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/LayerDaoJdbc.java @@ -158,6 +158,7 @@ public LayerDetail mapRow(ResultSet rs, int rowNum) throws SQLException { layer.services.addAll(Lists.newArrayList(rs.getString("str_services").split(","))); layer.timeout = rs.getInt("int_timeout"); layer.timeout_llu = rs.getInt("int_timeout_llu"); + layer.slotsRequired = rs.getInt("int_slots_required"); return layer; } }; @@ -241,7 +242,8 @@ public LayerInterface getLayer(String id) { + "int_dispatch_order, " + "str_tags, " + "str_type," + "int_cores_min, " + "int_cores_max, " + "b_threadable, " + "int_mem_min, " + "int_gpus_min, " + "int_gpus_max, " + "int_gpu_mem_min, " + "str_services, " + "int_timeout," - + "int_timeout_llu " + ") " + "VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)"; + + "int_timeout_llu, " + "int_slots_required " + ") " + + "VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)"; @Override public void insertLayerDetail(LayerDetail l) { @@ -250,7 +252,7 @@ public void insertLayerDetail(LayerDetail l) { l.chunkSize, l.dispatchOrder, StringUtils.join(l.tags, " | "), l.type.toString(), l.minimumCores, l.maximumCores, l.isThreadable, l.minimumMemory, l.minimumGpus, l.maximumGpus, l.minimumGpuMemory, StringUtils.join(l.services, ","), l.timeout, - l.timeout_llu); + l.timeout_llu, l.slotsRequired); } @Override @@ -571,6 +573,15 @@ public void updateTimeoutLLU(LayerInterface layer, int timeout_llu) { layer.getLayerId()); } + @Override + public void updateLayerSlotsRequired(LayerInterface layer, int slots) { + // Avoid negative numbers as they have the same meaning as zero + slots = Math.max(slots, 0); + + getJdbcTemplate().update("UPDATE layer SET int_slots_required=? WHERE pk_layer=?", slots, + layer.getLayerId()); + } + @Override public void enableMemoryOptimizer(LayerInterface layer, boolean value) { getJdbcTemplate().update("UPDATE layer SET b_optimize=? WHERE pk_layer=?", value, diff --git a/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/NestedWhiteboardDaoJdbc.java b/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/NestedWhiteboardDaoJdbc.java index 52f7afe19..c53e2015f 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/NestedWhiteboardDaoJdbc.java +++ b/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/NestedWhiteboardDaoJdbc.java @@ -283,9 +283,11 @@ private static final NestedJob mapResultSetToJob(ResultSet rs) throws SQLExcepti + "host.int_gpus, " + "host.int_gpus_idle, " + "host.int_gpu_mem, " + "host.int_gpu_mem_idle, " + "host.int_mem, " + "host.int_mem_idle, " + "host.str_lock_state, " + "host.str_tags, " + "host.b_comment, " - + "host.int_thread_mode, " + "host_stat.str_os, " + "host_stat.int_mem_total, " - + "host_stat.int_mem_free, " + "host_stat.int_swap_total, " - + "host_stat.int_swap_free, " + "host_stat.int_mcp_total, " + "host_stat.int_mcp_free, " + + "host.int_thread_mode, " + "host_stat.int_running_slots, " + + "host.int_concurrent_slots_limit, " + "host_stat.str_os, " + + "host_stat.int_mem_total, " + "host_stat.int_mem_free, " + + "host_stat.int_swap_total, " + "host_stat.int_swap_free, " + + "host_stat.int_mcp_total, " + "host_stat.int_mcp_free, " + "host_stat.int_gpu_mem_total, " + "host_stat.int_gpu_mem_free, " + "host_stat.int_load, " + "proc.pk_proc, " + "proc.int_cores_reserved AS proc_cores, " + "proc.int_gpus_reserved AS proc_gpus, " + "proc.int_mem_reserved AS proc_memory, " diff --git a/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/WhiteboardDaoJdbc.java b/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/WhiteboardDaoJdbc.java index d26f44360..23cd94b66 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/WhiteboardDaoJdbc.java +++ b/cuebot/src/main/java/com/imageworks/spcue/dao/postgres/WhiteboardDaoJdbc.java @@ -961,6 +961,8 @@ public static NestedHost.Builder mapNestedHostBuilder(ResultSet rs) throws SQLEx .setLockState(LockState.valueOf(SqlUtil.getString(rs, "str_lock_state"))) .setHasComment(rs.getBoolean("b_comment")) .setThreadMode(ThreadMode.values()[rs.getInt("int_thread_mode")]) + .setConcurrentSlotsLimit(rs.getInt("int_concurrent_slots_limit")) + .setRunningSlots(rs.getInt("int_running_slots")) .setOs(SqlUtil.getString(rs, "str_os")); String tags = SqlUtil.getString(rs, "str_tags"); @@ -998,6 +1000,8 @@ public static Host.Builder mapHostBuilder(ResultSet rs) throws SQLException { builder.setLockState(LockState.valueOf(SqlUtil.getString(rs, "str_lock_state"))); builder.setHasComment(rs.getBoolean("b_comment")); builder.setThreadMode(ThreadMode.values()[rs.getInt("int_thread_mode")]); + builder.setConcurrentSlotsLimit(rs.getInt("int_concurrent_slots_limit")); + builder.setRunningSlots(rs.getInt("int_running_slots")); builder.setOs(SqlUtil.getString(rs, "str_os")); String tags = SqlUtil.getString(rs, "str_tags"); @@ -1183,7 +1187,8 @@ public Layer mapRow(ResultSet rs, int rowNum) throws SQLException { Arrays.asList(SqlUtil.getString(rs, "str_limit_names").split(","))) .setMemoryOptimizerEnabled(rs.getBoolean("b_optimize")) .setTimeout(rs.getInt("int_timeout")) - .setTimeoutLlu(rs.getInt("int_timeout_llu")); + .setTimeoutLlu(rs.getInt("int_timeout_llu")) + .setSlotsRequired(rs.getInt("int_slots_required")); LayerStats.Builder statsBuilder = LayerStats.newBuilder() .setReservedCores(Convert.coreUnitsToCores(rs.getInt("int_cores"))) @@ -1716,13 +1721,14 @@ public Show mapRow(ResultSet rs, int rowNum) throws SQLException { + "host.int_cores_idle," + "host.int_mem," + "host.int_mem_idle," + "host.int_gpus," + "host.int_gpus_idle," + "host.int_gpu_mem," + "host.int_gpu_mem_idle," + "host.str_tags," + "host.str_lock_state," + "host.b_comment," - + "host.int_thread_mode," + "host_stat.str_os," + "host_stat.int_mem_total," - + "host_stat.int_mem_free," + "host_stat.int_swap_total," + "host_stat.int_swap_free," - + "host_stat.int_mcp_total," + "host_stat.int_mcp_free," + + "host.int_thread_mode," + "host.int_concurrent_slots_limit," + "host_stat.str_os," + + "host_stat.int_mem_total," + "host_stat.int_mem_free," + "host_stat.int_swap_total," + + "host_stat.int_swap_free," + "host_stat.int_mcp_total," + "host_stat.int_mcp_free," + "host_stat.int_gpu_mem_total," + "host_stat.int_gpu_mem_free," - + "host_stat.int_load, " + "alloc.str_name AS alloc_name " + "FROM " + "alloc," - + "facility, " + "host_stat," + "host " + "WHERE " + "host.pk_alloc = alloc.pk_alloc " - + "AND " + "facility.pk_facility = alloc.pk_facility " + "AND " + + "host_stat.int_load, " + "host_stat.int_running_slots, " + + "alloc.str_name AS alloc_name " + "FROM " + "alloc," + "facility, " + "host_stat," + + "host " + "WHERE " + "host.pk_alloc = alloc.pk_alloc " + "AND " + + "facility.pk_facility = alloc.pk_facility " + "AND " + "host.pk_host = host_stat.pk_host "; private static final String GET_DEPEND = "SELECT " + "depend.pk_depend, " + "depend.str_type, " diff --git a/cuebot/src/main/java/com/imageworks/spcue/dispatcher/DispatchSupportService.java b/cuebot/src/main/java/com/imageworks/spcue/dispatcher/DispatchSupportService.java index 3eb327e0f..8649b1f78 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/dispatcher/DispatchSupportService.java +++ b/cuebot/src/main/java/com/imageworks/spcue/dispatcher/DispatchSupportService.java @@ -392,7 +392,7 @@ public RunFrame prepareRqdRunFrame(VirtualProc proc, DispatchFrame frame) { .setNumCores(proc.coresReserved).setNumGpus(proc.gpusReserved) .setStartTime(System.currentTimeMillis()).setIgnoreNimby(proc.isLocalDispatch) .setOs(proc.os).setSoftMemoryLimit(frame.softMemoryLimit).setLokiUrl(frame.lokiURL) - .setHardMemoryLimit(frame.hardMemoryLimit) + .setHardMemoryLimit(frame.hardMemoryLimit).setSlotsRequired(proc.slotsRequired) .putAllEnvironment(jobDao.getEnvironment(frame)) .putAllEnvironment(layerDao.getLayerEnvironment(frame)).putEnvironment("CUE3", "1") .putEnvironment("CUE_THREADS", String.valueOf(threads)) diff --git a/cuebot/src/main/java/com/imageworks/spcue/dispatcher/HostReportHandler.java b/cuebot/src/main/java/com/imageworks/spcue/dispatcher/HostReportHandler.java index 4062ee13c..85602f52c 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/dispatcher/HostReportHandler.java +++ b/cuebot/src/main/java/com/imageworks/spcue/dispatcher/HostReportHandler.java @@ -185,13 +185,21 @@ public void handleHostReport(HostReport report, boolean isBoot) { DispatchHost host; RenderHost rhost = report.getHost(); + + /* + * Verify all the frames in the report are valid. Frames that are not valid are removed. + */ + List runningFrames = verifyRunningFrameInfo(report); + int host_running_slots = + runningFrames.stream().mapToInt(RunningFrameInfo::getSlotsRequired).sum(); + try { host = hostManager.findDispatchHost(rhost.getName()); hostManager.setHostStatistics(host, rhost.getTotalMem(), rhost.getFreeMem(), rhost.getTotalSwap(), rhost.getFreeSwap(), rhost.getTotalMcp(), rhost.getFreeMcp(), rhost.getTotalGpuMem(), rhost.getFreeGpuMem(), rhost.getLoad(), new Timestamp(rhost.getBootTime() * 1000l), - rhost.getAttributesMap().get("SP_OS")); + rhost.getAttributesMap().get("SP_OS"), host_running_slots); // Both logics are conflicting, only change hardware state if // there was no need for a tempDirStorage state change @@ -223,11 +231,6 @@ public void handleHostReport(HostReport report, boolean isBoot) { return; } - /* - * Verify all the frames in the report are valid. Frames that are not valid are removed. - */ - List runningFrames = verifyRunningFrameInfo(report); - /* * Updates memory usage for the proc, frames, jobs, and layers. And LLU time for the * frames. diff --git a/cuebot/src/main/java/com/imageworks/spcue/servant/ManageHost.java b/cuebot/src/main/java/com/imageworks/spcue/servant/ManageHost.java index 5732af62d..9ed1a25a2 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/servant/ManageHost.java +++ b/cuebot/src/main/java/com/imageworks/spcue/servant/ManageHost.java @@ -74,6 +74,8 @@ import com.imageworks.spcue.grpc.host.HostSetHardwareStateRequest; import com.imageworks.spcue.grpc.host.HostSetHardwareStateResponse; import com.imageworks.spcue.grpc.host.HostSetOsRequest; +import com.imageworks.spcue.grpc.host.HostSetConcurrentSlotsLimitRequest; +import com.imageworks.spcue.grpc.host.HostSetConcurrentSlotsLimitResponse; import com.imageworks.spcue.grpc.host.HostSetOsResponse; import com.imageworks.spcue.grpc.host.HostSetThreadModeRequest; import com.imageworks.spcue.grpc.host.HostSetThreadModeResponse; @@ -323,6 +325,15 @@ public void setOs(HostSetOsRequest request, responseObserver.onCompleted(); } + @Override + public void setConcurrentSlotsLimit(HostSetConcurrentSlotsLimitRequest request, + StreamObserver responseObserver) { + HostInterface host = getHostInterface(request.getHost()); + hostManager.setConcurrentSlotsLimit(host, request.getLimit()); + responseObserver.onNext(HostSetConcurrentSlotsLimitResponse.newBuilder().build()); + responseObserver.onCompleted(); + } + public HostManager getHostManager() { return hostManager; } diff --git a/cuebot/src/main/java/com/imageworks/spcue/servant/ManageLayer.java b/cuebot/src/main/java/com/imageworks/spcue/servant/ManageLayer.java index 0e19a2e43..21f64d9b3 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/servant/ManageLayer.java +++ b/cuebot/src/main/java/com/imageworks/spcue/servant/ManageLayer.java @@ -110,6 +110,8 @@ import com.imageworks.spcue.grpc.job.LayerSetTimeoutResponse; import com.imageworks.spcue.grpc.job.LayerSetTimeoutLLURequest; import com.imageworks.spcue.grpc.job.LayerSetTimeoutLLUResponse; +import com.imageworks.spcue.grpc.job.LayerSetSlotsRequiredRequest; +import com.imageworks.spcue.grpc.job.LayerSetSlotsRequiredResponse; import com.imageworks.spcue.grpc.job.LayerStaggerFramesRequest; import com.imageworks.spcue.grpc.job.LayerStaggerFramesResponse; import com.imageworks.spcue.grpc.limit.Limit; @@ -432,6 +434,15 @@ public void setTimeoutLLU(LayerSetTimeoutLLURequest request, } } + @Override + public void setSlotsRequired(LayerSetSlotsRequiredRequest request, + StreamObserver responseObserver) { + updateLayer(request.getLayer()); + jobManager.setLayerSlotsRequired(layer, request.getSlots()); + responseObserver.onNext(LayerSetSlotsRequiredResponse.newBuilder().build()); + responseObserver.onCompleted(); + } + @Override public void addLimit(LayerAddLimitRequest request, StreamObserver responseObserver) { diff --git a/cuebot/src/main/java/com/imageworks/spcue/servant/RqdReportStatic.java b/cuebot/src/main/java/com/imageworks/spcue/servant/RqdReportStatic.java index 2a30fb2cd..ef044834d 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/servant/RqdReportStatic.java +++ b/cuebot/src/main/java/com/imageworks/spcue/servant/RqdReportStatic.java @@ -12,11 +12,15 @@ import com.imageworks.spcue.grpc.report.RqdReportRunningFrameCompletionResponse; import com.imageworks.spcue.grpc.report.RqdReportStatusRequest; import com.imageworks.spcue.grpc.report.RqdReportStatusResponse; +import com.imageworks.spcue.grpc.report.RqdReportGetHostSlotsLimitRequest; +import com.imageworks.spcue.grpc.report.RqdReportGetHostSlotsLimitResponse; +import com.imageworks.spcue.service.HostManager; public class RqdReportStatic extends RqdReportInterfaceGrpc.RqdReportInterfaceImplBase { private FrameCompleteHandler frameCompleteHandler; private HostReportHandler hostReportHandler; + private HostManager hostManager; @SuppressWarnings("unused") @@ -44,6 +48,15 @@ public void reportStatus(RqdReportStatusRequest request, responseObserver.onCompleted(); } + @Override + public void getHostSlotsLimit(RqdReportGetHostSlotsLimitRequest request, + StreamObserver responseObserver) { + int slotsLimit = hostManager.getHostConcurrentSlotsLimit(request.getName()); + responseObserver.onNext( + RqdReportGetHostSlotsLimitResponse.newBuilder().setSlotsLimit(slotsLimit).build()); + responseObserver.onCompleted(); + } + public FrameCompleteHandler getFrameCompleteHandler() { return frameCompleteHandler; } @@ -59,4 +72,12 @@ public HostReportHandler getHostReportHandler() { public void setHostReportHandler(HostReportHandler hostReportHandler) { this.hostReportHandler = hostReportHandler; } + + public HostManager getHostManager() { + return hostManager; + } + + public void setHostManager(HostManager hostManager) { + this.hostManager = hostManager; + } } diff --git a/cuebot/src/main/java/com/imageworks/spcue/service/HostManager.java b/cuebot/src/main/java/com/imageworks/spcue/service/HostManager.java index 5913a7861..45c7eda98 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/service/HostManager.java +++ b/cuebot/src/main/java/com/imageworks/spcue/service/HostManager.java @@ -66,6 +66,22 @@ public interface HostManager { */ void setHostFreeTempDir(HostInterface host, Long freeTempDir); + /** + * Updates the concurrent procs limit of a host. + * + * @param host HostInterface + * @param limit int + */ + void setConcurrentSlotsLimit(HostInterface host, int limit); + + /** + * Gets the concurrent slots limit of a host by hostname. + * + * @param hostname String + * @return int the concurrent slots limit + */ + int getHostConcurrentSlotsLimit(String hostname); + DispatchHost createHost(HostReport report); DispatchHost createHost(RenderHost host); @@ -119,7 +135,7 @@ public interface HostManager { */ void setHostStatistics(HostInterface host, long totalMemory, long freeMemory, long totalSwap, long freeSwap, long totalMcp, long freeMcp, long totalGpuMemory, long freeGpuMemory, - int load, Timestamp bootTime, String os); + int load, Timestamp bootTime, String os, int runningProcs); void deleteHost(HostInterface host); diff --git a/cuebot/src/main/java/com/imageworks/spcue/service/HostManagerService.java b/cuebot/src/main/java/com/imageworks/spcue/service/HostManagerService.java index e7cee2248..f90524452 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/service/HostManagerService.java +++ b/cuebot/src/main/java/com/imageworks/spcue/service/HostManagerService.java @@ -93,6 +93,17 @@ public void setHostFreeTempDir(HostInterface host, Long freeTempDir) { hostDao.updateHostFreeTempDir(host, freeTempDir); } + @Override + public void setConcurrentSlotsLimit(HostInterface host, int limit) { + hostDao.updateConcurrentSlotsLimit(host, limit); + } + + @Override + @Transactional(propagation = Propagation.REQUIRED, readOnly = true) + public int getHostConcurrentSlotsLimit(String hostname) { + return hostDao.getHostConcurrentSlotsLimit(hostname); + } + public void rebootWhenIdle(HostInterface host) { try { hostDao.updateHostState(host, HardwareState.REBOOT_WHEN_IDLE); @@ -115,10 +126,10 @@ public void rebootNow(HostInterface host) { @Override public void setHostStatistics(HostInterface host, long totalMemory, long freeMemory, long totalSwap, long freeSwap, long totalMcp, long freeMcp, long totalGpuMemory, - long freeGpuMemory, int load, Timestamp bootTime, String os) { + long freeGpuMemory, int load, Timestamp bootTime, String os, int runningSlots) { hostDao.updateHostStats(host, totalMemory, freeMemory, totalSwap, freeSwap, totalMcp, - freeMcp, totalGpuMemory, freeGpuMemory, load, bootTime, os); + freeMcp, totalGpuMemory, freeGpuMemory, load, bootTime, os, runningSlots); } @Transactional(propagation = Propagation.SUPPORTS, readOnly = true) diff --git a/cuebot/src/main/java/com/imageworks/spcue/service/JobManager.java b/cuebot/src/main/java/com/imageworks/spcue/service/JobManager.java index 4641b8e82..51d6e9548 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/service/JobManager.java +++ b/cuebot/src/main/java/com/imageworks/spcue/service/JobManager.java @@ -455,6 +455,14 @@ public interface JobManager { */ void setLayerMinGpus(LayerInterface layer, int gpuUnits); + /** + * Sets the slots required for a layer. + * + * @param layer the layer to update + * @param slots the number of slots required + */ + void setLayerSlotsRequired(LayerInterface layer, int slots); + /** * Add a limit to the given layer. * diff --git a/cuebot/src/main/java/com/imageworks/spcue/service/JobManagerService.java b/cuebot/src/main/java/com/imageworks/spcue/service/JobManagerService.java index 03bc765b4..0904689fb 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/service/JobManagerService.java +++ b/cuebot/src/main/java/com/imageworks/spcue/service/JobManagerService.java @@ -456,6 +456,11 @@ public void setLayerMinGpus(LayerInterface layer, int gpu) { layerDao.updateLayerMinGpus(layer, gpu); } + @Override + public void setLayerSlotsRequired(LayerInterface layer, int slots) { + layerDao.updateLayerSlotsRequired(layer, slots); + } + @Override public void setLayerMaxGpus(LayerInterface layer, int gpu) { layerDao.updateLayerMaxGpus(layer, gpu); diff --git a/cuebot/src/main/java/com/imageworks/spcue/service/JobSpec.java b/cuebot/src/main/java/com/imageworks/spcue/service/JobSpec.java index d4ff4e6b5..bed7fa903 100644 --- a/cuebot/src/main/java/com/imageworks/spcue/service/JobSpec.java +++ b/cuebot/src/main/java/com/imageworks/spcue/service/JobSpec.java @@ -444,6 +444,10 @@ private void handleLayerTags(BuildableJob buildableJob, Element jobTag) { layer.timeout_llu = Integer.parseInt(layerTag.getChildTextTrim("timeout_llu")); } + if (layerTag.getChildTextTrim("slots_required") != null) { + layer.slotsRequired = Integer.parseInt(layerTag.getChildTextTrim("slots_required")); + } + /* * Handle the layer environment */ diff --git a/cuebot/src/main/resources/conf/ddl/postgres/migrations/V36__Add_layer_slots_required.sql b/cuebot/src/main/resources/conf/ddl/postgres/migrations/V36__Add_layer_slots_required.sql new file mode 100644 index 000000000..2172118dc --- /dev/null +++ b/cuebot/src/main/resources/conf/ddl/postgres/migrations/V36__Add_layer_slots_required.sql @@ -0,0 +1,4 @@ +-- Add a field to mark a layer as requiring at least a specific number of slots +-- <= 0 means slots are not required +alter table layer + add int_slots_required INT NOT NULL DEFAULT 0; diff --git a/cuebot/src/main/resources/conf/ddl/postgres/migrations/V38__Add_host_frame_slot_limit.sql b/cuebot/src/main/resources/conf/ddl/postgres/migrations/V38__Add_host_frame_slot_limit.sql new file mode 100644 index 000000000..8e315d778 --- /dev/null +++ b/cuebot/src/main/resources/conf/ddl/postgres/migrations/V38__Add_host_frame_slot_limit.sql @@ -0,0 +1,7 @@ +-- Add a field to limit the max amount of concurrent frames a host can run +-- -1 means no limit +alter table host + add int_concurrent_slots_limit INT NOT NULL DEFAULT -1; + +alter table host_stat + add int_running_slots INT NOT NULL DEFAULT 0; diff --git a/cuebot/src/main/resources/conf/spring/applicationContext-grpc.xml b/cuebot/src/main/resources/conf/spring/applicationContext-grpc.xml index 612aeaff5..f17040c50 100644 --- a/cuebot/src/main/resources/conf/spring/applicationContext-grpc.xml +++ b/cuebot/src/main/resources/conf/spring/applicationContext-grpc.xml @@ -196,6 +196,7 @@ + diff --git a/cuebot/src/main/resources/public/dtd/cjsl-1.16.dtd b/cuebot/src/main/resources/public/dtd/cjsl-1.16.dtd new file mode 100644 index 000000000..20ebf19c2 --- /dev/null +++ b/cuebot/src/main/resources/public/dtd/cjsl-1.16.dtd @@ -0,0 +1,106 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/cuebot/src/test/java/com/imageworks/spcue/test/dao/postgres/HostDaoTests.java b/cuebot/src/test/java/com/imageworks/spcue/test/dao/postgres/HostDaoTests.java index 796617fdc..caf755841 100644 --- a/cuebot/src/test/java/com/imageworks/spcue/test/dao/postgres/HostDaoTests.java +++ b/cuebot/src/test/java/com/imageworks/spcue/test/dao/postgres/HostDaoTests.java @@ -383,7 +383,8 @@ public void updateHostStats() { DispatchHost dispatchHost = hostDao.findDispatchHost(TEST_HOST); hostDao.updateHostStats(dispatchHost, CueUtil.GB8, CueUtil.GB8, CueUtil.GB8, CueUtil.GB8, - CueUtil.GB8, CueUtil.GB8, 1, 1, 100, new Timestamp(1247526000 * 1000l), "spinux1"); + CueUtil.GB8, CueUtil.GB8, 1, 1, 100, new Timestamp(1247526000 * 1000l), "spinux1", + 2); Map result = jdbcTemplate .queryForMap("SELECT * FROM host_stat WHERE pk_host=?", dispatchHost.getHostId()); diff --git a/cuebot/src/test/java/com/imageworks/spcue/test/dispatcher/HostReportHandlerTests.java b/cuebot/src/test/java/com/imageworks/spcue/test/dispatcher/HostReportHandlerTests.java index 3db3597fd..1a5f4141a 100644 --- a/cuebot/src/test/java/com/imageworks/spcue/test/dispatcher/HostReportHandlerTests.java +++ b/cuebot/src/test/java/com/imageworks/spcue/test/dispatcher/HostReportHandlerTests.java @@ -530,4 +530,180 @@ public void testMemoryAggressionMemoryWarning() { Math.max(memoryUsedProc3, layerBeforeIncrease.getMinimumMemory() + CueUtil.GB2), layer.getMinimumMemory()); } + + @Test + @Transactional + @Rollback(true) + public void testHostRunningSlotsWithNoFrames() { + CoreDetail cores = getCoreDetail(200, 200, 0, 0); + HostReport report = + HostReport.newBuilder().setHost(getRenderHost(hostname)).setCoreInfo(cores).build(); + + hostReportHandler.handleHostReport(report, false); + + DispatchHost host = getHost(hostname); + assertEquals(0, host.runningSlots); + } + + @Test + @Transactional + @Rollback(true) + public void testHostRunningSlotsWithSingleFrame() { + jobLauncher.testMode = true; + jobLauncher.launch(new File("src/test/resources/conf/jobspec/jobspec_simple.xml")); + + DispatchHost host = getHost(hostname); + List procs = dispatcher.dispatchHost(host); + assertEquals(1, procs.size()); + VirtualProc proc = procs.get(0); + + CoreDetail cores = getCoreDetail(200, 200, 0, 0); + int slotsRequired = 3; + + RunningFrameInfo info = RunningFrameInfo.newBuilder().setJobId(proc.getJobId()) + .setLayerId(proc.getLayerId()).setFrameId(proc.getFrameId()) + .setResourceId(proc.getProcId()).setSlotsRequired(slotsRequired).build(); + HostReport report = HostReport.newBuilder().setHost(getRenderHost(hostname)) + .setCoreInfo(cores).addFrames(info).build(); + + hostReportHandler.handleHostReport(report, false); + + host = getHost(hostname); + assertEquals(slotsRequired, host.runningSlots); + } + + @Test + @Transactional + @Rollback(true) + public void testHostRunningSlotsWithMultipleFrames() { + jobLauncher.testMode = true; + jobLauncher.launch(new File("src/test/resources/conf/jobspec/jobspec_multiple_frames.xml")); + + DispatchHost host = getHost(hostname); + List procs = dispatcher.dispatchHost(host); + assertEquals(3, procs.size()); + + VirtualProc proc1 = procs.get(0); + VirtualProc proc2 = procs.get(1); + VirtualProc proc3 = procs.get(2); + + CoreDetail cores = getCoreDetail(200, 200, 0, 0); + + int slotsRequired1 = 2; + int slotsRequired2 = 5; + int slotsRequired3 = 1; + int expectedTotalSlots = slotsRequired1 + slotsRequired2 + slotsRequired3; + + RunningFrameInfo info1 = RunningFrameInfo.newBuilder().setJobId(proc1.getJobId()) + .setLayerId(proc1.getLayerId()).setFrameId(proc1.getFrameId()) + .setResourceId(proc1.getProcId()).setSlotsRequired(slotsRequired1).build(); + + RunningFrameInfo info2 = RunningFrameInfo.newBuilder().setJobId(proc2.getJobId()) + .setLayerId(proc2.getLayerId()).setFrameId(proc2.getFrameId()) + .setResourceId(proc2.getProcId()).setSlotsRequired(slotsRequired2).build(); + + RunningFrameInfo info3 = RunningFrameInfo.newBuilder().setJobId(proc3.getJobId()) + .setLayerId(proc3.getLayerId()).setFrameId(proc3.getFrameId()) + .setResourceId(proc3.getProcId()).setSlotsRequired(slotsRequired3).build(); + + HostReport report = HostReport.newBuilder().setHost(getRenderHost(hostname)) + .setCoreInfo(cores).addAllFrames(Arrays.asList(info1, info2, info3)).build(); + + hostReportHandler.handleHostReport(report, false); + + host = getHost(hostname); + assertEquals(expectedTotalSlots, host.runningSlots); + } + + @Test + @Transactional + @Rollback(true) + public void testHostRunningSlotsWithZeroSlotsFrames() { + jobLauncher.testMode = true; + jobLauncher.launch(new File("src/test/resources/conf/jobspec/jobspec_multiple_frames.xml")); + + DispatchHost host = getHost(hostname); + List procs = dispatcher.dispatchHost(host); + assertEquals(3, procs.size()); + + VirtualProc proc1 = procs.get(0); + VirtualProc proc2 = procs.get(1); + + CoreDetail cores = getCoreDetail(200, 200, 0, 0); + + // Test with frames that have 0 slots required + RunningFrameInfo info1 = RunningFrameInfo.newBuilder().setJobId(proc1.getJobId()) + .setLayerId(proc1.getLayerId()).setFrameId(proc1.getFrameId()) + .setResourceId(proc1.getProcId()).setSlotsRequired(0).build(); + + RunningFrameInfo info2 = RunningFrameInfo.newBuilder().setJobId(proc2.getJobId()) + .setLayerId(proc2.getLayerId()).setFrameId(proc2.getFrameId()) + .setResourceId(proc2.getProcId()).setSlotsRequired(0).build(); + + HostReport report = HostReport.newBuilder().setHost(getRenderHost(hostname)) + .setCoreInfo(cores).addAllFrames(Arrays.asList(info1, info2)).build(); + + hostReportHandler.handleHostReport(report, false); + + host = getHost(hostname); + assertEquals(0, host.runningSlots); + } + + @Test + @Transactional + @Rollback(true) + public void testHostRunningSlotsUpdatesOnSubsequentReports() { + jobLauncher.testMode = true; + jobLauncher.launch(new File("src/test/resources/conf/jobspec/jobspec_multiple_frames.xml")); + + DispatchHost host = getHost(hostname); + List procs = dispatcher.dispatchHost(host); + assertEquals(3, procs.size()); + + VirtualProc proc1 = procs.get(0); + VirtualProc proc2 = procs.get(1); + VirtualProc proc3 = procs.get(2); + + CoreDetail cores = getCoreDetail(200, 200, 0, 0); + + // First report: 2 frames running with total 7 slots + RunningFrameInfo info1 = RunningFrameInfo.newBuilder().setJobId(proc1.getJobId()) + .setLayerId(proc1.getLayerId()).setFrameId(proc1.getFrameId()) + .setResourceId(proc1.getProcId()).setSlotsRequired(3).build(); + + RunningFrameInfo info2 = RunningFrameInfo.newBuilder().setJobId(proc2.getJobId()) + .setLayerId(proc2.getLayerId()).setFrameId(proc2.getFrameId()) + .setResourceId(proc2.getProcId()).setSlotsRequired(4).build(); + + HostReport report1 = HostReport.newBuilder().setHost(getRenderHost(hostname)) + .setCoreInfo(cores).addAllFrames(Arrays.asList(info1, info2)).build(); + + hostReportHandler.handleHostReport(report1, false); + host = getHost(hostname); + assertEquals(7, host.runningSlots); + + // Second report: 3 frames running with total 10 slots + RunningFrameInfo info3 = RunningFrameInfo.newBuilder().setJobId(proc3.getJobId()) + .setLayerId(proc3.getLayerId()).setFrameId(proc3.getFrameId()) + .setResourceId(proc3.getProcId()).setSlotsRequired(3).build(); + + HostReport report2 = HostReport.newBuilder().setHost(getRenderHost(hostname)) + .setCoreInfo(cores).addAllFrames(Arrays.asList(info1, info2, info3)).build(); + + hostReportHandler.handleHostReport(report2, false); + host = getHost(hostname); + assertEquals(10, host.runningSlots); + + // Third report: only 1 frame running with 2 slots + RunningFrameInfo info1Updated = RunningFrameInfo.newBuilder().setJobId(proc1.getJobId()) + .setLayerId(proc1.getLayerId()).setFrameId(proc1.getFrameId()) + .setResourceId(proc1.getProcId()).setSlotsRequired(2).build(); + + HostReport report3 = HostReport.newBuilder().setHost(getRenderHost(hostname)) + .setCoreInfo(cores).addFrames(info1Updated).build(); + + hostReportHandler.handleHostReport(report3, false); + host = getHost(hostname); + assertEquals(2, host.runningSlots); + } } diff --git a/cuebot/src/test/java/com/imageworks/spcue/test/service/JobSpecTests.java b/cuebot/src/test/java/com/imageworks/spcue/test/service/JobSpecTests.java index 69057bb83..533feb1c5 100644 --- a/cuebot/src/test/java/com/imageworks/spcue/test/service/JobSpecTests.java +++ b/cuebot/src/test/java/com/imageworks/spcue/test/service/JobSpecTests.java @@ -120,4 +120,26 @@ public void testParseMaxCoresAndMaxGpus() { assertEquals(job.maxGpusOverride, Integer.valueOf(42)); } + @Test + public void testParseSlotsRequired() { + String xml = readJobSpec("jobspec_1_16.xml"); + JobSpec spec = jobLauncher.parse(xml); + assertEquals(spec.getDoc().getDocType().getPublicID(), "SPI Cue Specification Language"); + assertEquals(spec.getDoc().getDocType().getSystemID(), + "http://localhost:8080/spcue/dtd/cjsl-1.16.dtd"); + assertEquals(spec.getJobs().size(), 1); + BuildableJob job = spec.getJobs().get(0); + assertEquals(job.getBuildableLayers().size(), 2); + + // First layer uses slot-based booking + LayerDetail slotBasedLayer = job.getBuildableLayers().get(0).layerDetail; + assertEquals(slotBasedLayer.name, "slot_based_layer"); + assertEquals(slotBasedLayer.slotsRequired, 4); + + // Second layer uses regular resource booking (default slots_required = 0) + LayerDetail regularLayer = job.getBuildableLayers().get(1).layerDetail; + assertEquals(regularLayer.name, "regular_layer"); + assertEquals(regularLayer.slotsRequired, 0); + } + } diff --git a/cuebot/src/test/resources/conf/dtd/cjsl-1.16.dtd b/cuebot/src/test/resources/conf/dtd/cjsl-1.16.dtd new file mode 100644 index 000000000..20ebf19c2 --- /dev/null +++ b/cuebot/src/test/resources/conf/dtd/cjsl-1.16.dtd @@ -0,0 +1,106 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/cuebot/src/test/resources/conf/jobspec/jobspec_1_16.xml b/cuebot/src/test/resources/conf/jobspec/jobspec_1_16.xml new file mode 100644 index 000000000..6e05da5bb --- /dev/null +++ b/cuebot/src/test/resources/conf/jobspec/jobspec_1_16.xml @@ -0,0 +1,58 @@ + + + + + local + testing + default + testuser + 9860 + + + False + 2 + 420 + 42 + False + + + + echo "Using slot-based booking" + 1-10 + 1 + 4 + + + shell + + + + echo "Using regular resource booking" + 1-5 + 1 + 100 + 2048 + + + shell + + + + + + diff --git a/cuegui/cuegui/HostMonitorTree.py b/cuegui/cuegui/HostMonitorTree.py index 207126214..073c9cdcb 100644 --- a/cuegui/cuegui/HostMonitorTree.py +++ b/cuegui/cuegui/HostMonitorTree.py @@ -159,6 +159,21 @@ def __init__(self, parent): data=lambda host: ",".join(host.data.tags), tip="The tags applied to the host.\n\n" "On a frame it is the name of the job.") + self.addColumn("Concurrent Slots", 50, id=23, + data=lambda host: \ + host.data.concurrent_slots_limit \ + if host.data.concurrent_slots_limit >= 0 \ + else "-", + tip="When >0 the host is configured to be slot based.\n" + "The host can only run this amount of slots at the same time " + "(Usually: 1 frame = 1 slot)\n\n" + "This host will only run layers with a slots_required field configured.") + self.addColumn("Available Slots", 50, id=24, + data=lambda host: ( + host.data.concurrent_slots_limit - host.data.running_slots + ) if host.data.concurrent_slots_limit >= 0 else "-", + tip="The number of slots still available on the host.\n" + "Computed as Concurrent Slots minus currently running slots.") self.hostSearch = opencue.search.HostSearch() @@ -290,6 +305,7 @@ def contextMenuEvent(self, e): self.__menuActions.hosts().addAction(menu, "removeTags") self.__menuActions.hosts().addAction(menu, "renameTag") self.__menuActions.hosts().addAction(menu, "changeAllocation") + self.__menuActions.hosts().addAction(menu, "setConcurrentSlotsLimit") self.__menuActions.hosts().addAction(menu, "delete") self.__menuActions.hosts().addAction(menu, "rebootWhenIdle") self.__menuActions.hosts().addAction(menu, "setRepair") diff --git a/cuegui/cuegui/LayerDialog.py b/cuegui/cuegui/LayerDialog.py index d95cb5367..10fc7545f 100644 --- a/cuegui/cuegui/LayerDialog.py +++ b/cuegui/cuegui/LayerDialog.py @@ -170,6 +170,12 @@ def __init__(self, layers, parent=None): self.__timeout_llu.setSuffix(" minutes") self.__timeout_llu.setSpecialValueText("No timeout") + # Slots Required + self.__slots_required = QtWidgets.QSpinBox(self) + self.__slots_required.setRange(0, int(self._cfg().get('max_cores', 16))) + self.__slots_required.setSingleStep(1) + self.__slots_required.setSpecialValueText("Not slot-based") + # Memory Optimizer self.__mem_opt = QtWidgets.QCheckBox() self.__mem_opt.setChecked(self.getMemoryOptSetting()) @@ -228,6 +234,7 @@ def __init__(self, layers, parent=None): self.__max_gpus.setValue(self.getMaxGpus()) self.__timeout.setValue(self.getTimeout()) self.__timeout_llu.setValue(self.getTimeoutLLU()) + self.__slots_required.setValue(self.getSlotsRequired()) QtWidgets.QVBoxLayout(self) @@ -272,6 +279,10 @@ def __init__(self, layers, parent=None): self.__timeout_llu, False), multiSelect)) + layout.addWidget(EnableableItem(LayerPropertiesItem("Slots Required:", + self.__slots_required, + False), + multiSelect)) layout.addStretch() self.__group.setLayout(layout) @@ -336,6 +347,8 @@ def apply(self): layer.setTimeout(self.__timeout.value()) if self.__timeout_llu.isEnabled(): layer.setTimeoutLLU(self.__timeout_llu.value()) + if self.__slots_required.isEnabled(): + layer.setSlotsRequired(self.__slots_required.value()) if self.__tags.isEnabled(): self.__tags.apply() if self.__limits.isEnabled(): @@ -421,6 +434,14 @@ def getMemoryOptSetting(self): break return result + def getSlotsRequired(self): + """Gets the layer slots required.""" + result = 0 + for layer in self.__layers: + if layer.data.slots_required > result: + result = layer.data.slots_required + return result + def __translateToMemSpinbox(self, value): self.__mem.spinner.setValue(float(value) / 1024.0) diff --git a/cuegui/cuegui/LayerMonitorTree.py b/cuegui/cuegui/LayerMonitorTree.py index 7c9a657e6..b49df189a 100644 --- a/cuegui/cuegui/LayerMonitorTree.py +++ b/cuegui/cuegui/LayerMonitorTree.py @@ -163,6 +163,13 @@ def __init__(self, parent): data=lambda layer: cuegui.Utils.secondsToHHHMM(layer.data.timeout_llu*60), sort=lambda layer: layer.data.timeout_llu, tip="Timeout for a frames\' LLU, Hours:Minutes") + self.addColumn("Slots Required", 65, id=23, + data=lambda layer: "-" \ + if layer.data.slots_required <= 0 \ + else str(layer.data.slots_required), + sort=lambda layer: layer.data.slots_required, + tip="Number of slots required per frame\n" + "(- means not slot-based)") cuegui.AbstractTreeWidget.AbstractTreeWidget.__init__(self, parent) # pylint: disable=no-member diff --git a/cuegui/cuegui/MenuActions.py b/cuegui/cuegui/MenuActions.py index fcdae5d41..e2c0e5fc9 100644 --- a/cuegui/cuegui/MenuActions.py +++ b/cuegui/cuegui/MenuActions.py @@ -1936,6 +1936,42 @@ def setThreadModeVariable(self, rpcObjects=None): host.setThreadMode("VARIABLE") self._update() + setConcurrentSlotsLimit_info = ["Update Slot Limit...", None, "configure"] + def setConcurrentSlotsLimit(self, rpcObjects=None): + """Set the concurrent slots limit for selected hosts.""" + hosts = self._getOnlyHostObjects(rpcObjects) + if not hosts: + return + + # Get current value from first selected host + current = hosts[0].concurrentSlotsLimit() if len(hosts) == 1 else 0 + + title = "Set Concurrent Slots Limit" + body = "Enter maximum concurrent slots \n(usually a frame consumes 1 slot, " \ + "the value can be configured on its layer's slot_required field)\n" \ + "When a limit is defined, booking will only allocate layers with " \ + "slots_required > 0 to be executed on this host. \n" \ + "Which means regular booking by cores/memory/gpu becomes disabled.\n\n" \ + "(0 for no limit, >0 for specific limit):" + + (value, choice) = QtWidgets.QInputDialog.getInt( + self._caller, + title, + body, + current, # current value + 0, # minimum value + 10000, # maximum value + 1, # step + ) + + if choice: + for host in hosts: + self.cuebotCall( + host.setConcurrentSlotsLimit, + "Set Concurrent Slots Limit on %s Failed" % host.data.name, + int(value), + ) + self._update() class ProcActions(AbstractActions): """Actions for procs.""" diff --git a/docs/_docs/concepts/glossary.md b/docs/_docs/concepts/glossary.md index 5e3f28e86..cee712ae2 100644 --- a/docs/_docs/concepts/glossary.md +++ b/docs/_docs/concepts/glossary.md @@ -4,7 +4,7 @@ nav_order: 13 parent: Concepts layout: default linkTitle: "Glossary" -date: 2019-02-22 +date: 2026-02-03 description: > Glossary of common OpenCue terms --- @@ -121,17 +121,17 @@ outline script to batch multiple job submissions into a single job, setting up dependencies and / or running in parallel. Outline scripts can submit almost any type of job to the cue, including Maya, Katana, or even shell commands. -## Proc - -A proc is a slot on a render *host* that has been carved out and isolated to -execute a *frame*. - ## OpenCue REST Gateway A production-ready HTTP service that translates REST API calls to gRPC communication with Cuebot. Enables web applications, scripts, and third-party tools to interact with OpenCue services through standard HTTP endpoints. +## Proc + +A proc is a unit of work on a render *host* that executes a *frame*. In +core-based booking, a proc represents reserved cores on a host. + ## PyCue The Python API library that provides programmatic access to OpenCue @@ -173,6 +173,17 @@ service by default. Services can be used to setup different requirements for different software jobs. For instance, a Maya render may need 6GB of memory vs a Nuke render may only need 2GB. +## Slot + +A slot is a scheduling unit used by slot-based booking. Slots are defined per +host and consumed per frame. + +## Slot-based booking + +A booking mode where hosts are scheduled by slots instead of cores, memory, or +GPU resources. It is enabled by setting a concurrent slots limit on hosts and a +`slots_required` value on layers. + ## Show A show is a group of related work to be done. *Jobs* submitted to OpenCue exist diff --git a/docs/_docs/concepts/slot-based-booking.md b/docs/_docs/concepts/slot-based-booking.md new file mode 100644 index 000000000..d23e17ccf --- /dev/null +++ b/docs/_docs/concepts/slot-based-booking.md @@ -0,0 +1,113 @@ +--- +title: "Slot-Based Booking" +nav_order: 21 +parent: Concepts +layout: default +linkTitle: "Slot-Based Booking" +date: 2026-02-03 +description: > + Limit per-host concurrency by slots instead of cores and memory. +--- + +# Slot-Based Booking + +### Limit per-host concurrency by slots instead of cores and memory + +--- + +Slot-based booking is an alternative scheduling mode that uses "slots" as the unit of capacity on a host. Instead of matching frames to available cores and memory, OpenCue can match frames to available slots when a layer requests them and a host is configured with a slot limit. + +## When to use slot-based booking + +- License-limited tools where you want a hard cap on concurrent frames per host +- I/O heavy tasks that should run fewer frames than the core count suggests +- Mixed hardware pools where cores and memory are not the best proxy for capacity + +## How it works + +- Hosts can define a `concurrent_slots_limit`. A positive value enables slot-based booking on that host. +- Layers can define `slots_required` per frame. A positive value marks the layer as slot-based. +- When a host has a slot limit, it only accepts layers with `slots_required > 0`. Core, memory, and GPU booking are disabled on that host. +- When a layer requests slots, the scheduler matches frames based on available slots. Use tags, services, and limits to constrain hardware or license requirements. + +> **Important** +> {: .callout .callout-warning} +> Setting a slot limit on a host changes how it is booked. The host will reject layers that do not set `slots_required`. + +## Configure hosts (concurrent slots limit) + +### CueCommander (CueGUI) + +1. Open **CueCommander** and go to **Monitor Hosts**. +1. Select one or more hosts. +1. Right-click and choose **Update Slot Limit...**. +1. Enter a value for the maximum concurrent slots. + +`0` disables slot-based booking for the host. A positive value enables slot-based booking and sets the maximum slots that can run concurrently. + +### PyCue + +```python +import opencue + +host = opencue.api.findHost("rendernode01") +host.setConcurrentSlotsLimit(8) +``` + +## Configure layers (slots required) + +### CueGUI layer dialog + +1. Open a layer's properties dialog. +1. Set **Slots Required**. + +`0` or a negative value means the layer is not slot-based. A positive value requests that many slots per frame. + +### PyOutline + +```python +import outline +import outline.modules.shell + +layer = outline.modules.shell.Shell( + "render", + command=["render", "-f", "#IFRAME#"], + range="1-100", + slots_required=2, +) + +# Or set it later: +layer.set_arg("slots_required", 2) +``` + +### PyCue (existing layer) + +```python +import opencue + +layer = opencue.api.findLayer("job_name", "render") +layer.setSlotsRequired(2) +``` + +### Job spec XML (CJS) + +```xml + + render -f #IFRAME# + 1-100 + 1 + 2 + +``` + +The `slots_required` element is available in the CJS 1.16 DTD. + +## Example: capacity math + +If a host has `concurrent_slots_limit = 8` and a layer uses `slots_required = 2`, that host can run 4 frames from that layer at the same time. + +## Operational notes + +- Slot-based booking is per host. Only hosts with a positive slot limit enforce slots. +- Use tags, services, and limits to target the right hosts when using slots. +- Keep slot values small and consistent. `1` is the common default unless you need heavier weighting. diff --git a/docs/_docs/user-guides/cuecommander-administration-guide.md b/docs/_docs/user-guides/cuecommander-administration-guide.md index b95a52ad1..4e78a3276 100644 --- a/docs/_docs/user-guides/cuecommander-administration-guide.md +++ b/docs/_docs/user-guides/cuecommander-administration-guide.md @@ -4,7 +4,7 @@ layout: default parent: User Guides nav_order: 38 linkTitle: "CueCommander Administration Guide" -date: 2025-01-13 +date: 2026-02-03 description: > Comprehensive guide to using CueCommander for OpenCue render farm administration --- @@ -301,6 +301,12 @@ Enables real-time monitoring and management of render hosts (nodes), including t - Identify stuck or long-running frames - Kill problematic processes +#### Slot-Based Booking Controls + +You can set a per-host slot limit from Monitor Hosts. Right-click one or more hosts and choose **Update Slot Limit...**. A value of `0` disables slot-based booking. A positive value enables slot-based booking and limits concurrent slots on the host. + +When a host has a slot limit, it only runs layers with `slots_required > 0`. See [Slot-Based Booking](/docs/concepts/slot-based-booking/) for details. + #### Common Use Cases - **Performance Troubleshooting**: Identify overloaded or underperforming hosts diff --git a/proto/src/host.proto b/proto/src/host.proto index 416e47320..0c700b98f 100644 --- a/proto/src/host.proto +++ b/proto/src/host.proto @@ -96,6 +96,9 @@ service HostInterface { // Changes the host's [ThreadMode] rpc SetThreadMode(HostSetThreadModeRequest) returns (HostSetThreadModeResponse); + // Set the maximum concurrent slots limit for the host. + rpc SetConcurrentSlotsLimit(HostSetConcurrentSlotsLimitRequest) returns (HostSetConcurrentSlotsLimitResponse); + // Unlocks the host for booking if the proc is in the Locked state. You cannot unlock a NimbyLocked proc. rpc Unlock(HostUnlockRequest) returns (HostUnlockResponse); } @@ -274,6 +277,12 @@ message Host { ThreadMode thread_mode = 27; float gpus = 28; float idle_gpus = 29; + + // Controls slot-based concurrent frame limiting on this host. + // 0: slot-based booking is disabled (default behavior using cores/memory/gpu) + // >0: slot-based booking is enabled with specified limit (only layers with slots_required > 0) + int32 concurrent_slots_limit = 30; + int32 running_slots = 31; } message HostSearchCriteria { @@ -321,6 +330,8 @@ message NestedHost { NestedProcSeq procs = 28; float gpus = 29; float idle_gpus = 30; + int32 concurrent_slots_limit = 31; + int32 running_slots = 32; } message NestedHostSeq { @@ -638,6 +649,14 @@ message HostSetThreadModeRequest { message HostSetThreadModeResponse {} // Empty +// SetConcurrentSlotsLimit +message HostSetConcurrentSlotsLimitRequest { + Host host = 1; + int32 limit = 2; +} + +message HostSetConcurrentSlotsLimitResponse {} // Empty + // Unlock message HostUnlockRequest { Host host = 1; diff --git a/proto/src/job.proto b/proto/src/job.proto index 13efc0a3c..2f0706bf8 100644 --- a/proto/src/job.proto +++ b/proto/src/job.proto @@ -392,6 +392,9 @@ service LayerInterface { // Set whether the LLU timeout for frames in the layer rpc SetTimeoutLLU(LayerSetTimeoutLLURequest) returns (LayerSetTimeoutLLUResponse); + // Set the number of slots required per frame for this layer + rpc SetSlotsRequired(LayerSetSlotsRequiredRequest) returns (LayerSetSlotsRequiredResponse); + // Staggers the specified frame range. rpc StaggerFrames(LayerStaggerFramesRequest) returns (LayerStaggerFramesResponse); } @@ -718,6 +721,8 @@ message Layer { float min_gpus = 20; float max_gpus = 21; string command = 22; + // Number of slots required per frame (<= 0 means not slot-based) + int32 slots_required = 23; } message LayerSeq { @@ -1799,6 +1804,14 @@ message LayerSetTimeoutLLURequest { message LayerSetTimeoutLLUResponse {} // Empty +// SetSlotsRequired +// +message LayerSetSlotsRequiredRequest { + Layer layer = 1; + int32 slots = 2; +} + +message LayerSetSlotsRequiredResponse {} // Empty // StaggerFrames message LayerStaggerFramesRequest { diff --git a/proto/src/report.proto b/proto/src/report.proto index 88d8d17d6..a582b0c59 100644 --- a/proto/src/report.proto +++ b/proto/src/report.proto @@ -23,6 +23,9 @@ service RqdReportInterface { // An incremental status report sent by RQD rpc ReportStatus(RqdReportStatusRequest) returns (RqdReportStatusResponse); + + // Get the host's slot limit + rpc GetHostSlotsLimit(RqdReportGetHostSlotsLimitRequest) returns (RqdReportGetHostSlotsLimitResponse); } @@ -109,6 +112,7 @@ message RunningFrameInfo { int64 used_swap_memory = 19; // kB int64 max_pss = 20; // kB int64 pss = 21; // kB + int32 slots_required = 22; }; message ChildrenProcStats { @@ -183,3 +187,12 @@ message RqdReportStatusRequest { HostReport host_report = 1; } message RqdReportStatusResponse {} // Empty + +// GetHostSlotsLimit +message RqdReportGetHostSlotsLimitRequest { + string name = 1; +} + +message RqdReportGetHostSlotsLimitResponse { + int32 slots_limit = 1; +} diff --git a/proto/src/rqd.proto b/proto/src/rqd.proto index 621a13212..2151ad9d6 100644 --- a/proto/src/rqd.proto +++ b/proto/src/rqd.proto @@ -126,6 +126,8 @@ message RunFrame { int64 hard_memory_limit = 27; int32 pid = 28; string loki_url = 29; + // Number of slots required per frame (<= 0 means not slot-based) + int32 slots_required = 30; } message RunFrameSeq { diff --git a/pycue/opencue/wrappers/host.py b/pycue/opencue/wrappers/host.py index c6fda1e87..de47308a5 100644 --- a/pycue/opencue/wrappers/host.py +++ b/pycue/opencue/wrappers/host.py @@ -128,6 +128,11 @@ def reboot(self): """Causes the host to kill all running frames and reboot the machine.""" self.stub.Reboot(host_pb2.HostRebootRequest(host=self.data), timeout=Cuebot.Timeout) + def setConcurrentSlotsLimit(self, limit): + """Set the concurrent slots limit for selected hosts.""" + self.stub.SetConcurrentSlotsLimit(host_pb2.HostSetConcurrentSlotsLimitRequest( + host=self.data, limit=limit), timeout=Cuebot.Timeout) + def addTags(self, tags): """Adds tags to a host. @@ -633,6 +638,15 @@ def os(self): """ return self.data.os + def concurrentSlotsLimit(self): + """Returns the limit of slots this host can run concurrently. + + :rtype: int + :return: the concurrent slots limit (0 = no limit) + """ + return self.data.concurrent_slots_limit + + class NestedHost(Host): """This class contains information and actions related to a nested host.""" diff --git a/pycue/opencue/wrappers/layer.py b/pycue/opencue/wrappers/layer.py index 3b5e9bcc2..a63f9f68d 100644 --- a/pycue/opencue/wrappers/layer.py +++ b/pycue/opencue/wrappers/layer.py @@ -185,6 +185,17 @@ def setMinMemory(self, memory): job_pb2.LayerSetMinMemoryRequest(layer=self.data, memory=memory), timeout=Cuebot.Timeout) + def setSlotsRequired(self, slots): + """Sets the number of slots required per frame for this layer. + + :type slots: int + :param slots: Number of slots required (=0 disables slot-based booking) + """ + return self.stub.SetSlotsRequired( + job_pb2.LayerSetSlotsRequiredRequest(layer=self.data, slots=slots), + timeout=Cuebot.Timeout, + ) + def setThreadable(self, threadable): """Sets the threadable field. @@ -457,6 +468,14 @@ def minMemory(self): """ return self.data.min_memory + def slotsRequired(self): + """Returns the number of slots required per frame. + + :rtype: int + :return: Number of slots required (<0 means not slot-based) + """ + return self.data.slots_required + def limits(self): """Returns the limit names for this layer. diff --git a/pycue/tests/wrappers/test_layer.py b/pycue/tests/wrappers/test_layer.py index 7adb7b11b..97c04d410 100644 --- a/pycue/tests/wrappers/test_layer.py +++ b/pycue/tests/wrappers/test_layer.py @@ -240,6 +240,20 @@ def testSetMinMemory(self, getStubMock): job_pb2.LayerSetMinMemoryRequest(layer=layer.data, memory=memory), timeout=mock.ANY) + def testSetSlotsRequired(self, getStubMock): + stubMock = mock.Mock() + stubMock.SetSlotsRequired.return_value = job_pb2.LayerSetSlotsRequiredResponse() + getStubMock.return_value = stubMock + + slots = 4 + layer = opencue.wrappers.layer.Layer(job_pb2.Layer(name=TEST_LAYER_NAME)) + layer.setSlotsRequired(slots) + + stubMock.SetSlotsRequired.assert_called_with( + job_pb2.LayerSetSlotsRequiredRequest(layer=layer.data, slots=slots), + timeout=mock.ANY, + ) + def testSetThreadable(self, getStubMock): stubMock = mock.Mock() stubMock.SetThreadable.return_value = job_pb2.LayerSetThreadableResponse() diff --git a/pyoutline/outline/backend/cue.py b/pyoutline/outline/backend/cue.py index ce88f3692..f6d6db991 100644 --- a/pyoutline/outline/backend/cue.py +++ b/pyoutline/outline/backend/cue.py @@ -374,6 +374,16 @@ def _serialize(launcher, use_pycuerun): else: _warning_spec_version(spec_version, "timeout_llu") + if layer.get_arg("slots_required"): + if spec_version >= Version("1.16"): + sub_element( + spec_layer, + "slots_required", + "%s" % (layer.get_arg("slots_required")), + ) + else: + _warning_spec_version(spec_version, "slots_required") + if os.environ.get("OL_TAG_OVERRIDE", False): sub_element(spec_layer, "tags", scrub_tags(os.environ["OL_TAG_OVERRIDE"])) diff --git a/pyoutline/outline/layer.py b/pyoutline/outline/layer.py index 856850ba7..d2eb23d31 100644 --- a/pyoutline/outline/layer.py +++ b/pyoutline/outline/layer.py @@ -15,25 +15,22 @@ """Base classes for all outline modules.""" -from __future__ import annotations -from __future__ import absolute_import -from __future__ import print_function -from __future__ import division +from __future__ import absolute_import, annotations, division, print_function +import logging import os import sys -import logging import tempfile from typing import ( - TypedDict, - List, - Optional, + Any, Callable, Dict, - Any, - Union, - Tuple, + List, + Optional, Set, + Tuple, + TypedDict, + Union, ) import FileSequence @@ -47,9 +44,9 @@ import outline.util if sys.version_info >= (3, 12): - from typing import override, Unpack + from typing import Unpack, override else: - from typing_extensions import override, Unpack + from typing_extensions import Unpack, override __all__ = [ "Layer", @@ -110,6 +107,8 @@ class _LayerArgs(TypedDict, total=False): # timeout_llu: Timeout for long last update in seconds # before considering a frame hung timeout_llu: int + # slots_required: Number of slots required per frame (<0 means not slot-based) + slots_required: int type: outline.constants.LayerType # The layer type (Render, Util, Post) @@ -1294,11 +1293,15 @@ class LayerPostProcess(Frame): the parent and the post process. """ - def __init__(self, creator: Layer, propigate: bool = True, **args: Unpack[_LayerArgs]) -> None: + def __init__( + self, creator: Layer, propigate: bool = True, **args: Unpack[_LayerArgs] + ) -> None: super().__init__(f"{creator.get_name()}_postprocess", **args) self.__creator = creator - self.depend_on(creator, outline.depend.DependType.LayerOnLayer, propigate=propigate) + self.depend_on( + creator, outline.depend.DependType.LayerOnLayer, propigate=propigate + ) self.set_type(outline.constants.LayerType.UTIL) diff --git a/rust/crates/dummy-cuebot/src/report_servant.rs b/rust/crates/dummy-cuebot/src/report_servant.rs index 1590dcb79..9eb290338 100644 --- a/rust/crates/dummy-cuebot/src/report_servant.rs +++ b/rust/crates/dummy-cuebot/src/report_servant.rs @@ -19,6 +19,9 @@ use opencue_proto::report::{ RqdReportRqdStartupResponse, RqdReportRunningFrameCompletionRequest, RqdReportRunningFrameCompletionResponse, RqdReportStatusRequest, RqdReportStatusResponse, }; +use opencue_proto::report::{ + RqdReportGetHostSlotsLimitRequest, RqdReportGetHostSlotsLimitResponse, +}; use tonic::transport::Server; use tonic::{async_trait, Request, Response, Status}; @@ -64,6 +67,23 @@ impl RqdReportInterface for ReportServant { Ok(Response::new(RqdReportStatusResponse {})) } + + /// Get the host's slot limit + async fn get_host_slots_limit( + &self, + request: tonic::Request, + ) -> std::result::Result, tonic::Status> + { + let name = request.into_inner().name; + println!( + "RqdReport: Received a get_host_slots_limit request with: {:?}", + name + ); + + Ok(Response::new(RqdReportGetHostSlotsLimitResponse { + slots_limit: -1, + })) + } } pub struct DummyCuebotServer {} diff --git a/rust/crates/dummy-cuebot/src/rqd_client.rs b/rust/crates/dummy-cuebot/src/rqd_client.rs index ab3ce01b4..bd9acdfd4 100644 --- a/rust/crates/dummy-cuebot/src/rqd_client.rs +++ b/rust/crates/dummy-cuebot/src/rqd_client.rs @@ -82,6 +82,8 @@ impl DummyRqdClient { #[allow(deprecated)] start_time: 0, + + slots_required: 0, }; let mut client = self.client.lock().await; diff --git a/rust/crates/rqd/src/frame/manager.rs b/rust/crates/rqd/src/frame/manager.rs index 811efebd1..d4cca5101 100644 --- a/rust/crates/rqd/src/frame/manager.rs +++ b/rust/crates/rqd/src/frame/manager.rs @@ -80,6 +80,8 @@ impl FrameManager { self.validate_grpc_frame(&run_frame)?; self.validate_machine_state(run_frame.ignore_nimby).await?; + let resource_id = run_frame.resource_id(); + // Create user if required. uid and gid ranges have already been verified let uid = match run_frame.uid_optional.as_ref().map(|o| match o { run_frame::UidOptional::Uid(v) => *v as u32, @@ -97,39 +99,13 @@ impl FrameManager { None => CONFIG.runner.default_uid, }; - // **Attention**: If an error happens between here and spawning a frame, the resources - // reserved need to be released. - - let num_cores = (run_frame.num_cores as u32).div_ceil(CONFIG.machine.core_multiplier); - - // Reserving cores will always yield a list of reserved thread_ids. If hyperthreading is off, - // the list should be ignored - let thread_ids = self - .machine - .reserve_cores(Either::Left(num_cores as usize), run_frame.resource_id()) - .await - .map_err(|err| { - FrameManagerError::Aborted(format!( - "Not launching, failed to reserve cpu resources {:?}", - err - )) - })?; // Although num_gpus is not required on a frame, the field is not optional on the proto // layer. =0 means None, !=0 means Some let gpu_list = match run_frame.num_gpus { 0 => None, _ => { + // TODO: Release GPUs in case of error when GPU support gets implemented let reserved_res = self.machine.reserve_gpus(run_frame.num_gpus as u32).await; - if reserved_res.is_err() { - // Release cores reserved on the last step - if let Err(err) = self.machine.release_cores(&run_frame.resource_id()).await { - warn!( - "Failed to release cores reserved for {} during gpu reservation failure. {}", - &run_frame.resource_id(), - err - ) - }; - } Some(reserved_res.map_err(|err| { FrameManagerError::Aborted(format!( "Not launching, insufficient resources {:?}", @@ -145,27 +121,90 @@ impl FrameManager { .environment .get("CUE_THREADABLE") .is_some_and(|v| v == "1"); - // Ignore the list of allocated threads if hyperthreading is off - let thread_ids = hyperthreaded.then_some(thread_ids); - let resource_id = run_frame.resource_id(); - let running_frame = Arc::new(RunningFrame::init( - run_frame, - uid, - CONFIG.runner.clone(), - thread_ids, - gpu_list, - self.machine.get_host_name().await, - )); + let slot_based_booking = self.machine.is_slot_configured().await; + // Keep track of reserved slots, if any + let mut reserved_slots = 0; + + let running_frame = match slot_based_booking { + // Core based booking + false => { + // **Attention**: If an error happens between here and spawning a frame, the resources + // reserved need to be released. + let num_cores = + (run_frame.num_cores as u32).div_ceil(CONFIG.machine.core_multiplier); + + // Reserving cores will always yield a list of reserved thread_ids. If hyperthreading is off, + // the list should be ignored + let thread_ids = self + .machine + .reserve_cores(Either::Left(num_cores as usize), run_frame.resource_id()) + .await + .map_err(|err| { + FrameManagerError::Aborted(format!( + "Not launching, failed to reserve cpu resources {:?}", + err + )) + })?; + // Ignore the list of allocated threads if hyperthreading is off + let thread_ids = hyperthreaded.then_some(thread_ids); + + Arc::new(RunningFrame::init( + run_frame, + uid, + CONFIG.runner.clone(), + thread_ids, + gpu_list, + self.machine.get_host_name(), + )) + } + // Slot based booking + true => { + reserved_slots = if run_frame.slots_required > 0 { + run_frame.slots_required as u32 + } else { + Err(FrameManagerError::InvalidArgument( + "Core based frame cannot be launched on a slot configured host".to_string(), + ))? + }; + self.machine + .reserve_slots(reserved_slots) + .await + .map_err(|err| { + FrameManagerError::Aborted(format!( + "Not launching, failed to reserve {:} slots {:?}", + run_frame.slots_required, err + )) + })?; + + Arc::new(RunningFrame::init( + run_frame, + uid, + CONFIG.runner.clone(), + // Disable taskset to avoid binding this frame to specific threads + None, + gpu_list, + self.machine.get_host_name(), + )) + } + }; if cfg!(feature = "containerized_frames") && CONFIG.runner.run_on_docker { #[cfg(feature = "containerized_frames")] self.spawn_docker_frame(running_frame, false); } else if self.spawn_running_frame(running_frame, false).is_err() { - // Release cores reserved if spawning the frame failed - if let Err(err) = self.machine.release_cores(&resource_id).await { + let release_res = if slot_based_booking { + // Release slots reserved if spawning the frame failed + self.machine.release_slots(reserved_slots).await + } else { + // Release cores reserved if spawning the frame failed + self.machine.release_cores(&resource_id).await + }; + + // Log failure to release + if let Err(err) = release_res { warn!( - "Failed to release cores reserved for {} during spawn failure. {}", + "Failed to release resources reserved for {} during spawn failure. {}", &resource_id, err ); } @@ -207,43 +246,75 @@ impl FrameManager { }) .collect(); let mut errors = Vec::new(); + let slot_based_booking = self.machine.is_slot_configured().await; + for path in snapshot_dir { let running_frame = RunningFrame::from_snapshot(&path, CONFIG.runner.clone()) .await .map(Arc::new); match running_frame { Ok(running_frame) => { - // Update reservations. If a thread_ids list exists, the frame was booked using affinity - if let Err(err) = match &running_frame.thread_ids { - Some(thread_ids) => { - self.machine - .reserve_cores( - Either::Right(thread_ids.clone()), - running_frame.request.resource_id(), - ) - .await + let resource_id = running_frame.request.resource_id(); + let mut reserved_slots = 0; + + // Update reservations based on booking mode + if let Err(err) = match slot_based_booking { + // Core-based booking: If a thread_ids list exists, the frame was booked using affinity + false => { + match &running_frame.thread_ids { + Some(thread_ids) => { + self.machine + .reserve_cores( + Either::Right(thread_ids.clone()), + running_frame.request.resource_id(), + ) + .await + } + None => { + let num_cores = (running_frame.request.num_cores as u32) + .div_ceil(CONFIG.machine.core_multiplier); + self.machine + .reserve_cores( + Either::Left(num_cores as usize), + running_frame.request.resource_id(), + ) + .await + } + } + // Ignore reserved threads as they are no longer necessary + .map(|_| ()) } - None => { - let num_cores = (running_frame.request.num_cores as u32) - .div_ceil(CONFIG.machine.core_multiplier); - self.machine - .reserve_cores( - Either::Left(num_cores as usize), - running_frame.request.resource_id(), - ) - .await + // Slot-based booking + true => { + reserved_slots = if running_frame.request.slots_required > 0 { + running_frame.request.slots_required as u32 + } else { + errors.push(format!( + "Core based frame {} cannot be recovered on a slot configured host", + resource_id + )); + continue; + }; + self.machine.reserve_slots(reserved_slots).await } } { errors.push(err.to_string()); } - let resource_id = running_frame.request.resource_id(); if CONFIG.runner.run_on_docker { todo!("Recovering frames when running on docker is not yet supported") } else if self.spawn_running_frame(running_frame, true).is_err() { - if let Err(err) = self.machine.release_cores(&resource_id).await { + let release_res = if slot_based_booking { + // Release slots reserved if spawning the frame failed + self.machine.release_slots(reserved_slots).await + } else { + self.machine.release_cores(&resource_id).await + }; + + // Failed to release + if let Err(err) = release_res { warn!( - "Failed to release cores reserved for {} during recover spawn error. {}", + "Failed to release resources reserved for {} during recover spawn error. {}", &resource_id, err ); } diff --git a/rust/crates/rqd/src/frame/running_frame.rs b/rust/crates/rqd/src/frame/running_frame.rs index f1949de6e..5ce2ffee3 100644 --- a/rust/crates/rqd/src/frame/running_frame.rs +++ b/rust/crates/rqd/src/frame/running_frame.rs @@ -1387,6 +1387,7 @@ mod tests { #[allow(deprecated)] start_time: 0, + slots_required: 0, }, uid, config, diff --git a/rust/crates/rqd/src/report/report_client.rs b/rust/crates/rqd/src/report/report_client.rs index 31995e894..092cb9944 100644 --- a/rust/crates/rqd/src/report/report_client.rs +++ b/rust/crates/rqd/src/report/report_client.rs @@ -183,6 +183,7 @@ pub trait ReportInterface { run_time: u32, ) -> Result<()>; async fn send_host_report(&self, host_report: pb::HostReport) -> Result<()>; + async fn get_host_slots_limit(&self, name: String) -> Result>; } #[async_trait] @@ -242,4 +243,23 @@ impl ReportInterface for ReportClient { .into_diagnostic() .and(Ok(())) } + + async fn get_host_slots_limit(&self, name: String) -> Result> { + let request = pb::RqdReportGetHostSlotsLimitRequest { name }; + let slots_limit = self + .get_client() + .await? + .get_host_slots_limit(request) + .await + .into_diagnostic()? + .into_inner() + .slots_limit; + + // Host with limit <= 0 are running on core based booking mode, so they don't have a limit + if slots_limit > 0 { + Ok(Some(slots_limit as u32)) + } else { + Ok(None) + } + } } diff --git a/rust/crates/rqd/src/system/machine.rs b/rust/crates/rqd/src/system/machine.rs index ad24d5ac7..8950b2300 100644 --- a/rust/crates/rqd/src/system/machine.rs +++ b/rust/crates/rqd/src/system/machine.rs @@ -79,6 +79,11 @@ pub struct MachineMonitor { pub core_manager: Arc>, pub running_frames_cache: Arc, last_host_state: Arc>>, + // Host name is only written once at the beginning of start. After that it is only read. + // This makes it safe to have a sync lock to give the object mutability (it can't be + // initialized at init) but avoid unecessary awaits + host_name: std::sync::RwLock>, + slot_state: RwLock>, interrupt: Mutex>>, reboot_when_idle: Mutex, #[cfg(feature = "nimby")] @@ -87,6 +92,11 @@ pub struct MachineMonitor { nimby_state: RwLock, } +struct SlotState { + slot_limit: u32, + slots_consumed: u32, +} + static MACHINE_MONITOR: OnceCell> = OnceCell::const_new(); pub async fn instance() -> Result> { @@ -154,6 +164,7 @@ impl MachineMonitor { system_manager: Mutex::new(system_manager), running_frames_cache: RunningFrameCache::init(), last_host_state: Arc::new(RwLock::new(None)), + host_name: std::sync::RwLock::new(None), interrupt: Mutex::new(None), reboot_when_idle: Mutex::new(false), #[cfg(feature = "nimby")] @@ -161,6 +172,7 @@ impl MachineMonitor { #[cfg(feature = "nimby")] nimby_state: RwLock::new(LockState::Open), core_manager, + slot_state: RwLock::new(None), }) } @@ -178,6 +190,14 @@ impl MachineMonitor { core_manager.get_core_info_report(self.maching_config.core_multiplier) }; + // Write host_name to the object + { + self.host_name + .write() + .unwrap_or_else(|p| p.into_inner()) + .replace(host_state.name.clone()); + } + self.last_host_state .write() .await @@ -217,6 +237,7 @@ impl MachineMonitor { _ = interval.tick() => { self.collect_and_send_host_report().await?; self.check_reboot_flag().await; + self.check_host_state_on_server().await; #[cfg(feature = "nimby")] if let Some(nimby) = &*self.nimby { @@ -338,6 +359,29 @@ impl MachineMonitor { } } + async fn check_host_state_on_server(&self) { + let client = self.report_client.clone(); + + if let Some(slot_limit) = client + .get_host_slots_limit(self.get_host_name()) + .await + .ok() + .flatten() + { + let mut current_state = self.slot_state.write().await; + let slots_consumed = current_state + .as_ref() + .map(|s| s.slots_consumed) + .unwrap_or(0); + + // Replace limit but keep consumed count + current_state.replace(SlotState { + slot_limit, + slots_consumed, + }); + } + } + async fn monitor_running_frames(&self) -> Result<()> { let mut finished_frames: Vec> = Vec::new(); let mut running_frames: Vec<(Arc, RunningState)> = Vec::new(); @@ -504,7 +548,19 @@ impl MachineMonitor { let frame_report = frame.clone_into_running_frame_info(); info!("Sending frame complete report: {}", frame); - if let Err(err) = self.release_cores(&frame.request.resource_id()).await { + // Either release slots or cores, depending on whether it was configured with slots + if frame.request.slots_required > 0 { + if let Err(err) = self + .release_slots(frame.request.slots_required as u32) + .await + { + warn!( + "Failed to release slots reserved by {}: {}", + frame.request.resource_id(), + err + ); + }; + } else if let Err(err) = self.release_cores(&frame.request.resource_id()).await { warn!( "Failed to release cores reserved by {}: {}", frame.request.resource_id(), @@ -572,6 +628,7 @@ pub trait Machine { async fn hardware_state(&self) -> Option; async fn memory_usage(&self) -> Option<(u32, u64)>; async fn nimby_locked(&self) -> bool; + async fn is_slot_configured(&self) -> bool; /// Reserve CPU cores for a resource /// @@ -593,6 +650,40 @@ pub trait Machine { resource_id: Uuid, ) -> Result, ReservationError>; + /// Reserve slot units for a resource + /// + /// # Arguments + /// + /// * `requested_slots` - Number of slots to reserve + /// + /// # Returns + /// + /// Returns `Ok(())` if the slots were successfully reserved + /// + /// # Errors + /// + /// Returns `ReservationError` if: + /// * There are not enough available slots (`NotEnoughResourcesAvailable`) + /// * Slot reservation is not configured on this machine (`InvalidSlotReservationRequest`) + async fn reserve_slots(&self, requested_slots: u32) -> Result<(), ReservationError>; + + /// Release slot units previously reserved by a resource + /// + /// # Arguments + /// + /// * `requested_slots` - Number of slots to release + /// + /// # Returns + /// + /// Returns `Ok(())` if the slots were successfully released + /// + /// # Errors + /// + /// Returns `ReservationError` if: + /// * Attempting to release more slots than are currently consumed (`NotEnoughResourcesAvailable`) + /// * Slot reservation is not configured on this machine (`InvalidSlotReservationRequest`) + async fn release_slots(&self, requested_slots: u32) -> Result<(), ReservationError>; + /// Release CPU cores previously reserved by a resource /// /// # Arguments @@ -632,7 +723,17 @@ pub trait Machine { /// The user ID (uid) of the created or existing user async fn create_user_if_unexisting(&self, username: &str, uid: u32, gid: u32) -> Result; - async fn get_host_name(&self) -> String; + /// Returns the hostname of this machine + /// + /// The hostname is determined during the initial startup report and remains + /// constant throughout the machine's lifecycle. If the hostname hasn't been + /// initialized yet (which shouldn't happen in normal operation), returns + /// "noname" as a fallback. + /// + /// # Returns + /// + /// The machine's hostname as a String + fn get_host_name(&self) -> String; /// Send a signal to kill a process /// @@ -698,6 +799,10 @@ impl Machine for MachineMonitor { .unwrap_or(false) } + async fn is_slot_configured(&self) -> bool { + self.slot_state.read().await.as_ref().is_some() + } + async fn reserve_cores( &self, request: Either>, @@ -711,6 +816,38 @@ impl Machine for MachineMonitor { } } + async fn reserve_slots(&self, requested_slots: u32) -> Result<(), ReservationError> { + let mut slot_state = self.slot_state.write().await; + + match slot_state.as_mut() { + Some(slot_state) => { + if slot_state.slots_consumed + requested_slots <= slot_state.slot_limit { + slot_state.slots_consumed += requested_slots; + Ok(()) + } else { + Err(ReservationError::NotEnoughResourcesAvailable) + } + } + None => Err(ReservationError::InvalidSlotReservationRequest), + } + } + + async fn release_slots(&self, released_slots: u32) -> Result<(), ReservationError> { + let mut slot_state = self.slot_state.write().await; + + match slot_state.as_mut() { + Some(slot_state) => { + if released_slots <= slot_state.slots_consumed { + slot_state.slots_consumed -= released_slots; + Ok(()) + } else { + Err(ReservationError::NotEnoughResourcesAvailable) + } + } + None => Err(ReservationError::InvalidSlotReservationRequest), + } + } + async fn release_cores(&self, resource_id: &Uuid) -> Result<(), ReservationError> { let mut core_manager = self.core_manager.write().await; core_manager.release_cores(resource_id).map(|_| ()) @@ -725,11 +862,11 @@ impl Machine for MachineMonitor { system.create_user_if_unexisting(username, uid, gid) } - async fn get_host_name(&self) -> String { - let lock = self.last_host_state.read().await; - - lock.as_ref() - .map(|h| h.name.clone()) + fn get_host_name(&self) -> String { + self.host_name + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clone() .unwrap_or("noname".to_string()) } diff --git a/rust/crates/rqd/src/system/manager.rs b/rust/crates/rqd/src/system/manager.rs index 5ba276c19..fa15448ad 100644 --- a/rust/crates/rqd/src/system/manager.rs +++ b/rust/crates/rqd/src/system/manager.rs @@ -70,6 +70,9 @@ pub enum ReservationError { #[error("Could not find core owner of this thread id")] CoreNotFoundForThread(Vec), + + #[error("Slot reservation requested when the host is configured to core based booking")] + InvalidSlotReservationRequest, } /// Represents attributes on a machine that should never change without restarting the diff --git a/rust/crates/rqd/src/system/oom.rs b/rust/crates/rqd/src/system/oom.rs index 2ca82f69e..d5a91a820 100644 --- a/rust/crates/rqd/src/system/oom.rs +++ b/rust/crates/rqd/src/system/oom.rs @@ -239,6 +239,7 @@ mod tests { log_dir_file: "".to_string(), #[allow(deprecated)] start_time: 0, + slots_required: 0, }, 1000, config, diff --git a/rust/crates/scheduler/src/config/mod.rs b/rust/crates/scheduler/src/config/mod.rs index 908d61453..dcd499e0e 100644 --- a/rust/crates/scheduler/src/config/mod.rs +++ b/rust/crates/scheduler/src/config/mod.rs @@ -218,6 +218,8 @@ pub struct HostCacheConfig { #[serde(with = "humantime_serde")] pub host_staleness_threshold: Duration, pub update_stat_on_book: bool, + pub slot_min_cores: u32, + pub slot_min_memory: ByteSize, } impl Default for HostCacheConfig { @@ -232,6 +234,8 @@ impl Default for HostCacheConfig { concurrent_fetch_permit: 4, host_staleness_threshold: Duration::from_secs(2 * 60), // 2 minutes update_stat_on_book: false, + slot_min_cores: 1, + slot_min_memory: ByteSize::mib(257), } } } diff --git a/rust/crates/scheduler/src/dao/frame_dao.rs b/rust/crates/scheduler/src/dao/frame_dao.rs index 9459abd22..55beb6a96 100644 --- a/rust/crates/scheduler/src/dao/frame_dao.rs +++ b/rust/crates/scheduler/src/dao/frame_dao.rs @@ -67,6 +67,7 @@ pub struct DispatchFrameModel { pub b_threadable: bool, pub int_gpus_min: i64, pub int_gpu_mem_min: i64, + pub int_slots_required: i32, // On Cuebot these fields come from constants, maybe replicate these constants here // pub int_soft_memory_limit: i64, // pub int_hard_memory_limit: i64, @@ -155,6 +156,7 @@ impl From for DispatchFrame { version: val.int_version as u32, updated_at, env, + slots_required: val.int_slots_required as u32, } } } diff --git a/rust/crates/scheduler/src/dao/host_dao.rs b/rust/crates/scheduler/src/dao/host_dao.rs index 4d6d37fef..b4ea96d67 100644 --- a/rust/crates/scheduler/src/dao/host_dao.rs +++ b/rust/crates/scheduler/src/dao/host_dao.rs @@ -71,6 +71,8 @@ pub struct HostModel { // Number of cores available at the subscription of the show this host has been queried on int_alloc_available_cores: i64, ts_ping: DateTime, + int_concurrent_slots_limit: i64, + int_running_slots: i64, } impl From for Host { @@ -105,6 +107,9 @@ impl From for Host { alloc_id: parse_uuid(&val.pk_alloc), alloc_name: val.str_alloc_name, last_updated: val.ts_ping, + concurrent_slots_limit: (val.int_concurrent_slots_limit > 0) + .then_some(val.int_concurrent_slots_limit as u32), + running_slots_count: val.int_running_slots as u32, } } } @@ -130,7 +135,9 @@ SELECT DISTINCT s.int_burst - s.int_cores as int_alloc_available_cores, a.pk_alloc, a.str_name as str_alloc_name, - hs.ts_ping + hs.ts_ping, + h.int_concurrent_slots_limit, + hs.int_running_slots FROM host h INNER JOIN host_stat hs ON h.pk_host = hs.pk_host INNER JOIN alloc a ON h.pk_alloc = a.pk_alloc @@ -152,13 +159,15 @@ WHERE pk_host = $5 RETURNING int_cores_idle, int_mem_idle, int_gpus_idle, int_gpu_mem_idle, NOW() "#; -// This update is meant for testing environments where rqd is not constantly reporting -// host reports to Cuebot to get host_stats properly updated. +// ATTENTION: This update is meant for testing environments where rqd is not constantly reporting +// host reports to Cuebot to get host_stats properly updated. This is turned of by default and +// can be turned on by `host_cache.update_stat_on_book=true` static UPDATE_HOST_STAT: &str = r#" UPDATE host_stat SET int_mem_free = int_mem_free - $1, - int_gpu_mem_free = int_gpu_mem_free - $2 -WHERE pk_host = $3 + int_gpu_mem_free = int_gpu_mem_free - $2, + int_running_slots = int_running_slots + $3 +WHERE pk_host = $4 "#; static UPDATE_SUBSCRIPTION: &str = r#" @@ -347,6 +356,7 @@ impl HostDao { sqlx::query(UPDATE_HOST_STAT) .bind((virtual_proc.memory_reserved.as_u64() / KB) as i64) .bind(virtual_proc.gpu_memory_reserved.as_u64() as i64) + .bind(virtual_proc.slots_required as i64) .bind(host_id.to_string()) .execute(&mut **transaction) .await diff --git a/rust/crates/scheduler/src/dao/layer_dao.rs b/rust/crates/scheduler/src/dao/layer_dao.rs index 62c007c9e..4817d1365 100644 --- a/rust/crates/scheduler/src/dao/layer_dao.rs +++ b/rust/crates/scheduler/src/dao/layer_dao.rs @@ -56,6 +56,7 @@ pub struct DispatchLayerModel { pub b_threadable: bool, pub int_gpus_min: i64, pub int_gpu_mem_min: i64, + pub int_slots_required: i32, pub str_tags: String, } @@ -79,6 +80,7 @@ pub struct LayerWithFramesModel { pub b_threadable: bool, pub int_gpus_min: i64, pub int_gpu_mem_min: i64, + pub int_slots_required: i32, pub str_tags: String, pub job_env: Json>, pub layer_env: Json>, @@ -134,13 +136,19 @@ impl DispatchLayer { ), mem_min: ByteSize::kb(layer.int_mem_min as u64), threadable: layer.b_threadable, - gpus_min: layer - .int_gpus_min - .try_into() - .expect("gpus_min should fit on a i32"), + gpus_min: CoreSize( + layer + .int_gpus_min + .try_into() + .expect("gpus_min should fit on a i32"), + ), gpu_mem_min: ByteSize::kb(layer.int_gpu_mem_min as u64), tags: layer.str_tags.split(" | ").map(|t| t.to_string()).collect(), frames: frames.into_iter().map(|f| f.into()).collect(), + slots_required: layer + .int_slots_required + .try_into() + .expect("int_slots_required should fit on a i32"), } } } @@ -220,6 +228,7 @@ SELECT DISTINCT l.b_threadable, l.int_gpus_min, l.int_gpu_mem_min, + l.int_slots_required, l.str_tags, je.job_env, le.layer_env, @@ -359,6 +368,7 @@ impl LayerDao { int_gpus_min: model.int_gpus_min, int_gpu_mem_min: model.int_gpu_mem_min, str_tags: model.str_tags.clone(), + int_slots_required: model.int_slots_required, }; // Extract frame data (if present) @@ -382,6 +392,7 @@ impl LayerDao { str_job_name: model.job_name.clone(), int_min_cores: model.int_layer_cores_min.unwrap_or(100), // default core multiplier int_mem_min: model.int_mem_min_frame.unwrap_or(0), + int_slots_required: model.int_slots_required, b_threadable: model.b_threadable, int_gpus_min: model.int_gpus_min_frame.unwrap_or(0), int_gpu_mem_min: model.int_gpu_mem_min_frame.unwrap_or(0), diff --git a/rust/crates/scheduler/src/host_cache/actor.rs b/rust/crates/scheduler/src/host_cache/actor.rs index 40d098850..fbeff087d 100644 --- a/rust/crates/scheduler/src/host_cache/actor.rs +++ b/rust/crates/scheduler/src/host_cache/actor.rs @@ -12,7 +12,6 @@ use actix::{Actor, ActorFutureExt, AsyncContext, Handler, ResponseActFuture, WrapFuture}; -use bytesize::ByteSize; use itertools::Itertools; use miette::IntoDiagnostic; use scc::{hash_map::OccupiedEntry, HashMap, HashSet}; @@ -35,7 +34,7 @@ use crate::{ config::CONFIG, dao::HostDao, host_cache::{messages::*, store, *}, - models::{CoreSize, Host}, + models::{Host, ResourceRequest}, }; #[derive(Clone)] @@ -105,8 +104,7 @@ where facility_id, show_id, tags, - cores, - memory, + resource_request, validation, } = msg; @@ -115,7 +113,7 @@ where Box::pin( async move { let out = service - .check_out(facility_id, show_id, tags, cores, memory, validation) + .check_out(facility_id, show_id, tags, resource_request, validation) .await; if let Ok(host) = &out { debug!("Checked out {}", host.1); @@ -208,8 +206,7 @@ impl HostCacheService { facility_id: Uuid, show_id: Uuid, tags: Vec, - cores: CoreSize, - memory: ByteSize, + resource_request: ResourceRequest, validation: F, ) -> Result where @@ -235,9 +232,9 @@ impl HostCacheService { // fight for the same rows. .read_async(&cache_key, |_, cached_group| { if !cached_group.expired() { + // Checkout host from a group cached_group - // Checkout host from a group - .check_out(cores, memory, validation) + .check_out(resource_request, validation) .map(|host| (cache_key.clone(), host.clone())) .ok() } else { @@ -260,7 +257,7 @@ impl HostCacheService { .map_err(|err| HostCacheError::FailedToQueryHostCache(err.to_string()))?; let checked_out_host = group // Checkout host from a group - .check_out(cores, memory, validation) + .check_out(resource_request, validation) .map(|host| CheckedOutHost(cache_key.clone(), host.clone())); if let Ok(checked_out_host) = checked_out_host { diff --git a/rust/crates/scheduler/src/host_cache/cache.rs b/rust/crates/scheduler/src/host_cache/cache.rs index d41252d34..5e5d63a0f 100644 --- a/rust/crates/scheduler/src/host_cache/cache.rs +++ b/rust/crates/scheduler/src/host_cache/cache.rs @@ -46,7 +46,7 @@ use uuid::Uuid; use crate::{ config::{HostBookingStrategy, CONFIG}, host_cache::{store::HOST_STORE, HostCacheError, HostId}, - models::{CoreSize, Host}, + models::{CoreSize, Host, ResourceRequest}, }; type CoreKey = u32; @@ -144,8 +144,7 @@ impl HostCache { /// /// # Arguments /// - /// * `cores` - Minimum number of cores required - /// * `memory` - Minimum memory required + /// * `resource_request` - The resource requirements (cores and memory, GPU, or unit) /// * `validation` - Function to validate additional host requirements /// /// # Returns @@ -154,8 +153,7 @@ impl HostCache { /// * `Err(HostCacheError)` - No suitable host available pub fn check_out( &self, - cores: CoreSize, - memory: ByteSize, + resource_request: ResourceRequest, validation: F, ) -> Result where @@ -163,9 +161,24 @@ impl HostCache { { self.ping_query(); - let host = self - .remove_host(cores, memory, validation) - .ok_or(HostCacheError::NoCandidateAvailable)?; + let host = match resource_request { + ResourceRequest::CoresAndMemory { cores, memory } => self + .remove_host(cores, memory, 1, validation) + .ok_or(HostCacheError::NoCandidateAvailable)?, + ResourceRequest::Gpu { cores, memory } => { + todo!("GPU host search is not yet implemented. Request: {cores}, {memory}") + } + ResourceRequest::Slots(slots) => self + // Request a host with minimum requirements as the remove logic already accounts for + // limiting slots + .remove_host( + CoreSize(CONFIG.host_cache.slot_min_cores as i32), + CONFIG.host_cache.slot_min_memory, + slots, + validation, + ) + .ok_or(HostCacheError::NoCandidateAvailable)?, + }; Ok(host) } @@ -186,7 +199,13 @@ impl HostCache { /// /// * `Some(Host)` - Host that meets all requirements /// * `None` - No suitable host found - fn remove_host(&self, cores: CoreSize, memory: ByteSize, validation: F) -> Option + fn remove_host( + &self, + cores: CoreSize, + memory: ByteSize, + slots: u32, + validation: F, + ) -> Option where F: Fn(&Host) -> bool, { @@ -200,6 +219,8 @@ impl HostCache { // Check memory and core requirements just in case host.idle_memory >= memory && host.idle_cores >= cores && + host.running_slots_count + slots + <= host.concurrent_slots_limit.unwrap_or(u32::MAX) && // Ensure we're not retrying the same host as last attempts !failed_candidates.borrow().contains(&host.id) }; @@ -366,6 +387,8 @@ mod tests { alloc_id: Uuid::new_v4(), alloc_name: "test".to_string(), last_updated: Utc::now(), + concurrent_slots_limit: None, + running_slots_count: 0, } } @@ -460,8 +483,10 @@ mod tests { cache.check_in(host, false); let result = cache.check_out( - CoreSize(2), - ByteSize::gb(4), + ResourceRequest::CoresAndMemory { + cores: CoreSize(2), + memory: ByteSize::gb(4), + }, |_| true, // Always validate true ); @@ -485,7 +510,13 @@ mod tests { fn test_checkout_no_candidate_available() { let cache = HostCache::default(); - let result = cache.check_out(CoreSize(4), ByteSize::gb(8), |_| true); + let result = cache.check_out( + ResourceRequest::CoresAndMemory { + cores: CoreSize(4), + memory: ByteSize::gb(8), + }, + |_| true, + ); assert!(result.is_err()); assert!(matches!(result, Err(HostCacheError::NoCandidateAvailable))); @@ -500,8 +531,10 @@ mod tests { cache.check_in(host, false); let result = cache.check_out( - CoreSize(4), // Request more cores than available - ByteSize::gb(4), + ResourceRequest::CoresAndMemory { + cores: CoreSize(4), // Request more cores than available + memory: ByteSize::gb(4), + }, |_| true, ); @@ -517,8 +550,10 @@ mod tests { cache.check_in(host, false); let result = cache.check_out( - CoreSize(2), - ByteSize::gb(8), // Request more memory than available + ResourceRequest::CoresAndMemory { + cores: CoreSize(2), + memory: ByteSize::gb(8), // Request more memory than available + }, |_| true, ); @@ -534,8 +569,10 @@ mod tests { cache.check_in(host, false); let result = cache.check_out( - CoreSize(2), - ByteSize::gb(4), + ResourceRequest::CoresAndMemory { + cores: CoreSize(2), + memory: ByteSize::gb(4), + }, |_| false, // Always fail validation ); @@ -551,11 +588,23 @@ mod tests { cache.check_in(host, false); // First checkout should succeed - let result1 = cache.check_out(CoreSize(2), ByteSize::gb(4), |_| true); + let result1 = cache.check_out( + ResourceRequest::CoresAndMemory { + cores: CoreSize(2), + memory: ByteSize::gb(4), + }, + |_| true, + ); assert!(result1.is_ok()); // Second checkout should fail because host is already checked out - let result2 = cache.check_out(CoreSize(2), ByteSize::gb(4), |_| true); + let result2 = cache.check_out( + ResourceRequest::CoresAndMemory { + cores: CoreSize(2), + memory: ByteSize::gb(4), + }, + |_| true, + ); assert!(result2.is_err()); } @@ -568,7 +617,13 @@ mod tests { cache.check_in(host.clone(), false); // Checkout the host - let mut checked_host = assert_ok!(cache.check_out(CoreSize(2), ByteSize::gb(4), |_| true)); + let mut checked_host = assert_ok!(cache.check_out( + ResourceRequest::CoresAndMemory { + cores: CoreSize(2), + memory: ByteSize::gb(4), + }, + |_| true + )); assert_eq!(checked_host.idle_cores.value(), 4); // Reduce the number of cores and checkin to ensure cache is updated @@ -576,8 +631,20 @@ mod tests { // Check it back in cache.check_in(checked_host, false); - assert_err!(cache.check_out(CoreSize(2), ByteSize::gb(4), |_| true)); - assert_ok!(cache.check_out(CoreSize(1), ByteSize::gb(4), |_| true)); + assert_err!(cache.check_out( + ResourceRequest::CoresAndMemory { + cores: CoreSize(2), + memory: ByteSize::gb(4), + }, + |_| true + )); + assert_ok!(cache.check_out( + ResourceRequest::CoresAndMemory { + cores: CoreSize(1), + memory: ByteSize::gb(4), + }, + |_| true + )); } #[test] @@ -599,7 +666,13 @@ mod tests { cache.check_in(host3, false); // Request 3 cores, 6GB - should get host2 (4 cores, 8GB) or host3 (8 cores, 16GB) - let result = cache.check_out(CoreSize(3), ByteSize::gb(6), |_| true); + let result = cache.check_out( + ResourceRequest::CoresAndMemory { + cores: CoreSize(3), + memory: ByteSize::gb(6), + }, + |_| true, + ); assert!(result.is_ok()); let chosen_host = result.unwrap(); @@ -643,11 +716,23 @@ mod tests { cache.check_in(host2, false); // First checkout should succeed - let result1 = cache.check_out(CoreSize(2), ByteSize::gb(4), |_| true); + let result1 = cache.check_out( + ResourceRequest::CoresAndMemory { + cores: CoreSize(2), + memory: ByteSize::gb(4), + }, + |_| true, + ); assert!(result1.is_ok()); // Second checkout should also succeed (different host) - let result2 = cache.check_out(CoreSize(2), ByteSize::gb(4), |_| true); + let result2 = cache.check_out( + ResourceRequest::CoresAndMemory { + cores: CoreSize(2), + memory: ByteSize::gb(4), + }, + |_| true, + ); assert!(result2.is_ok()); // The hosts should be different diff --git a/rust/crates/scheduler/src/host_cache/messages.rs b/rust/crates/scheduler/src/host_cache/messages.rs index bb25ced14..94ba38af1 100644 --- a/rust/crates/scheduler/src/host_cache/messages.rs +++ b/rust/crates/scheduler/src/host_cache/messages.rs @@ -12,14 +12,13 @@ use actix::{Message, MessageResponse}; -use bytesize::ByteSize; use miette::Result; use uuid::Uuid; use crate::{ cluster_key::{ClusterKey, Tag}, host_cache::HostCacheError, - models::{CoreSize, Host}, + models::{Host, ResourceRequest}, }; /// Response containing a checked-out host and its associated cluster key. @@ -66,8 +65,7 @@ where pub facility_id: Uuid, pub show_id: Uuid, pub tags: Vec, - pub cores: CoreSize, - pub memory: ByteSize, + pub resource_request: ResourceRequest, pub validation: F, } diff --git a/rust/crates/scheduler/src/host_cache/store.rs b/rust/crates/scheduler/src/host_cache/store.rs index 5509e23e1..aec35764b 100644 --- a/rust/crates/scheduler/src/host_cache/store.rs +++ b/rust/crates/scheduler/src/host_cache/store.rs @@ -400,6 +400,8 @@ mod tests { alloc_id: Uuid::new_v4(), alloc_name: "test".to_string(), last_updated, + concurrent_slots_limit: None, + running_slots_count: 0, } } diff --git a/rust/crates/scheduler/src/models/frame.rs b/rust/crates/scheduler/src/models/frame.rs index b83aaa2e0..0172bba63 100644 --- a/rust/crates/scheduler/src/models/frame.rs +++ b/rust/crates/scheduler/src/models/frame.rs @@ -51,6 +51,7 @@ pub struct DispatchFrame { pub min_gpus: u32, pub min_gpu_memory: ByteSize, pub min_memory: ByteSize, + pub slots_required: u32, // On Cuebot these fields come from constants, maybe replicate these constants here // pub int_soft_memory_limit: i64, // pub int_hard_memory_limit: i64, diff --git a/rust/crates/scheduler/src/models/host.rs b/rust/crates/scheduler/src/models/host.rs index 38af88726..d50143441 100644 --- a/rust/crates/scheduler/src/models/host.rs +++ b/rust/crates/scheduler/src/models/host.rs @@ -36,6 +36,8 @@ pub struct Host { pub(crate) alloc_id: Uuid, pub(crate) alloc_name: String, pub(crate) last_updated: DateTime, + pub(crate) concurrent_slots_limit: Option, + pub(crate) running_slots_count: u32, } impl Host { @@ -75,6 +77,7 @@ impl Host { alloc_available_cores: CoreSize, alloc_id: Uuid, alloc_name: String, + concurrent_frames_limit: Option, ) -> Self { Self { id, @@ -91,6 +94,8 @@ impl Host { alloc_id, alloc_name, last_updated: Local::now().with_timezone(&Utc), + concurrent_slots_limit: concurrent_frames_limit, + running_slots_count: 0, } } } diff --git a/rust/crates/scheduler/src/models/layer.rs b/rust/crates/scheduler/src/models/layer.rs index 10f795fa8..2cd27dadf 100644 --- a/rust/crates/scheduler/src/models/layer.rs +++ b/rust/crates/scheduler/src/models/layer.rs @@ -31,8 +31,9 @@ pub struct DispatchLayer { pub cores_min: CoreSize, pub mem_min: ByteSize, pub threadable: bool, - pub gpus_min: i32, + pub gpus_min: CoreSize, pub gpu_mem_min: ByteSize, + pub slots_required: u32, pub tags: HashSet, pub frames: Vec, } @@ -49,6 +50,17 @@ impl fmt::Display for DispatchLayer { } } +/// Describes what resources are required to run a frame from this layer +#[derive(Clone, Copy)] +pub enum ResourceRequest { + /// Request a machine with at least this amount of cores and memory idle + CoresAndMemory { cores: CoreSize, memory: ByteSize }, + /// Request a machine with this amount of gpu cores idle + Gpu { cores: CoreSize, memory: ByteSize }, + /// Request a machine with this amount of frame slots available + Slots(u32), +} + impl DispatchLayer { /// Removes frames with matching IDs from this layer's frame list. /// @@ -61,4 +73,20 @@ impl DispatchLayer { pub fn drain_frames(&mut self, frame_ids: Vec) { self.frames.retain(|f| !frame_ids.contains(&f.id)) } + + pub fn resource_request(&self) -> ResourceRequest { + if self.slots_required > 0 { + ResourceRequest::Slots(self.slots_required) + } else if self.gpus_min.value() > 0 { + ResourceRequest::Gpu { + cores: self.gpus_min, + memory: self.gpu_mem_min, + } + } else { + ResourceRequest::CoresAndMemory { + cores: self.cores_min, + memory: self.mem_min, + } + } + } } diff --git a/rust/crates/scheduler/src/models/mod.rs b/rust/crates/scheduler/src/models/mod.rs index 9f61fd86d..50767d267 100644 --- a/rust/crates/scheduler/src/models/mod.rs +++ b/rust/crates/scheduler/src/models/mod.rs @@ -22,7 +22,7 @@ pub use core_size::{CoreSize, CoreSizeWithMultiplier}; pub use frame::DispatchFrame; pub use host::Host; pub use job::DispatchJob; -pub use layer::DispatchLayer; +pub use layer::{DispatchLayer, ResourceRequest}; pub use subscription::{Allocation, Subscription}; pub use virtual_proc::VirtualProc; diff --git a/rust/crates/scheduler/src/models/virtual_proc.rs b/rust/crates/scheduler/src/models/virtual_proc.rs index 5ae6dc190..82363434d 100644 --- a/rust/crates/scheduler/src/models/virtual_proc.rs +++ b/rust/crates/scheduler/src/models/virtual_proc.rs @@ -35,6 +35,7 @@ pub struct VirtualProc { pub os: String, pub is_local_dispatch: bool, pub frame: DispatchFrame, + pub slots_required: u32, } impl Display for VirtualProc { diff --git a/rust/crates/scheduler/src/pipeline/dispatcher/actor.rs b/rust/crates/scheduler/src/pipeline/dispatcher/actor.rs index bcebd4e7d..048e8ee10 100644 --- a/rust/crates/scheduler/src/pipeline/dispatcher/actor.rs +++ b/rust/crates/scheduler/src/pipeline/dispatcher/actor.rs @@ -721,6 +721,15 @@ impl RqdDispatcherService { )))? } + if let Some(concurrent_slots_limit) = host.concurrent_slots_limit { + if concurrent_slots_limit < host.running_slots_count + frame.slots_required { + Err(VirtualProcError::HostResourcesExtinguished(format!( + "Not enough slots available: {} slots taken, requested {} slots", + host.running_slots_count, frame.slots_required + )))? + } + } + let memory_reserved = frame.min_memory; let gpus_reserved = frame.min_gpus; let gpu_memory_reserved = frame.min_gpu_memory; @@ -733,6 +742,9 @@ impl RqdDispatcherService { ByteSize(host.idle_gpu_memory.as_u64() - gpu_memory_reserved.as_u64()); // Field will be overwritten with database values as soon as the changes are committed host.last_updated = Utc::now(); + if host.concurrent_slots_limit.is_some() { + host.running_slots_count += frame.slots_required; + } Ok(( VirtualProc { @@ -751,6 +763,7 @@ impl RqdDispatcherService { is_local_dispatch: false, frame: frame.clone(), host_name: host.name.clone(), + slots_required: frame.slots_required, }, host, )) @@ -970,6 +983,7 @@ impl RqdDispatcherService { log_file: "deprecated".to_string(), #[allow(deprecated)] log_dir_file: "deprecated".to_string(), + slots_required: proc.slots_required as i32, }; Ok(run_frame) @@ -1043,6 +1057,7 @@ mod tests { CoreSize(4), Uuid::new_v4(), "test-alloc".to_string(), + None, ) } @@ -1078,6 +1093,7 @@ mod tests { version: 1, updated_at: SystemTime::now(), env: HashMap::new(), + slots_required: 0, } } @@ -1176,6 +1192,7 @@ mod tests { CoreSize(4), Uuid::new_v4(), "test-alloc".to_string(), + None, ); let mut frame = create_test_dispatch_frame(); @@ -1212,6 +1229,7 @@ mod tests { CoreSize(4), Uuid::new_v4(), "test-alloc".to_string(), + None, ); let mut frame = create_test_dispatch_frame(); @@ -1248,6 +1266,7 @@ mod tests { CoreSize(8), Uuid::new_v4(), "test-alloc".to_string(), + None, ); let mut frame = create_test_dispatch_frame(); @@ -1283,6 +1302,7 @@ mod tests { CoreSize(8), Uuid::new_v4(), "test-alloc".to_string(), + None, ); let mut frame = create_test_dispatch_frame(); @@ -1460,6 +1480,7 @@ mod tests { is_local_dispatch: false, frame, host_name: "somehost".to_string(), + slots_required: 0, }; let result = RqdDispatcherService::prepare_rqd_run_frame(&virtual_proc); @@ -1536,6 +1557,7 @@ mod tests { is_local_dispatch: false, frame, host_name: "somehost".to_string(), + slots_required: 0, }; let result = RqdDispatcherService::prepare_rqd_run_frame(&virtual_proc); @@ -1572,9 +1594,267 @@ mod tests { is_local_dispatch: false, frame, host_name: "somehost".to_string(), + slots_required: 0, }; let result = RqdDispatcherService::prepare_rqd_run_frame(&virtual_proc); assert!(result.is_err()); } + + // ── Slot-based scheduling tests ────────────────────────────────────────── + + #[tokio::test] + async fn test_consume_host_virtual_resources_slots_required_propagated_to_virtual_proc() { + // When frame.slots_required != 0, the VirtualProc should carry the same value + let mut frame = create_test_dispatch_frame(); + frame.slots_required = 2; + + let mut host = create_test_host(); + host.concurrent_slots_limit = Some(4); + + let result = RqdDispatcherService::consume_host_virtual_resources( + &frame, + &host, + ByteSize::gib(1), + ) + .await; + + assert!(result.is_ok()); + let (virtual_proc, _updated_host) = result.unwrap(); + assert_eq!(virtual_proc.slots_required, 2); + } + + #[tokio::test] + async fn test_consume_host_virtual_resources_with_slots_deducts_cores_and_memory() { + // Slot-based frames still deduct core and memory resources from the host + let mut frame = create_test_dispatch_frame(); + frame.slots_required = 2; + frame.min_cores = CoreSize(1); + frame.min_memory = ByteSize::gib(2); + frame.threadable = false; // predictable core reservation + + let mut host = create_test_host(); + host.concurrent_slots_limit = Some(4); + let initial_idle_memory = host.idle_memory; + let initial_idle_cores = host.idle_cores; + + let result = RqdDispatcherService::consume_host_virtual_resources( + &frame, + &host, + ByteSize::gib(1), + ) + .await; + + assert!(result.is_ok()); + let (virtual_proc, updated_host) = result.unwrap(); + + assert_eq!( + updated_host.idle_memory.as_u64(), + initial_idle_memory.as_u64() - frame.min_memory.as_u64() + ); + assert!(updated_host.idle_cores < initial_idle_cores); + assert_eq!(virtual_proc.memory_reserved, frame.min_memory); + assert_eq!(virtual_proc.slots_required, 2); + } + + #[tokio::test] + async fn test_consume_host_virtual_resources_with_slots_updates_running_slots_count() { + // When concurrent_slots_limit is set, running_slots_count should be incremented + let mut frame = create_test_dispatch_frame(); + frame.slots_required = 2; + frame.threadable = false; + + let mut host = create_test_host(); + host.concurrent_slots_limit = Some(4); + host.running_slots_count = 0; + + let result = RqdDispatcherService::consume_host_virtual_resources( + &frame, + &host, + ByteSize::gib(1), + ) + .await; + + assert!(result.is_ok()); + let (_virtual_proc, updated_host) = result.unwrap(); + assert_eq!(updated_host.running_slots_count, 2); + } + + #[tokio::test] + async fn test_consume_host_virtual_resources_with_slots_fails_on_insufficient_cores() { + // Slot-based frames still fail when the host cannot satisfy core requirements + let mut frame = create_test_dispatch_frame(); + frame.slots_required = 2; + frame.min_cores = CoreSize(100); // exceeds host capacity + + let mut host = create_test_host(); + host.concurrent_slots_limit = Some(4); + + let result = RqdDispatcherService::consume_host_virtual_resources( + &frame, + &host, + ByteSize::gib(1), + ) + .await; + + assert!(result.is_err()); + match result { + Err(VirtualProcError::HostResourcesExtinguished(msg)) => { + assert!(msg.contains("Not enough cores")); + } + _ => panic!("Expected HostResourcesExtinguished error for insufficient cores"), + } + } + + #[tokio::test] + async fn test_consume_host_virtual_resources_with_slots_fails_on_insufficient_memory() { + // Slot-based frames still fail when the host cannot satisfy memory requirements + let mut frame = create_test_dispatch_frame(); + frame.slots_required = 2; + frame.min_memory = ByteSize::gib(64); // exceeds host capacity + + let mut host = create_test_host(); + host.concurrent_slots_limit = Some(4); + + let result = RqdDispatcherService::consume_host_virtual_resources( + &frame, + &host, + ByteSize::gib(1), + ) + .await; + + assert!(result.is_err()); + match result { + Err(VirtualProcError::HostResourcesExtinguished(msg)) => { + assert!(msg.contains("Not enough memory")); + } + _ => panic!("Expected HostResourcesExtinguished error for insufficient memory"), + } + } + + #[tokio::test] + async fn test_consume_host_virtual_resources_with_slots_fails_when_slots_exhausted() { + // When running_slots_count + slots_required exceeds concurrent_slots_limit, dispatch fails + let mut frame = create_test_dispatch_frame(); + frame.slots_required = 3; + frame.threadable = false; + + let mut host = create_test_host(); + host.concurrent_slots_limit = Some(4); + host.running_slots_count = 2; // 2 + 3 > 4 + + let result = RqdDispatcherService::consume_host_virtual_resources( + &frame, + &host, + ByteSize::gib(1), + ) + .await; + + assert!(result.is_err()); + match result { + Err(VirtualProcError::HostResourcesExtinguished(msg)) => { + assert!(msg.contains("Not enough slots")); + } + _ => panic!("Expected HostResourcesExtinguished error for insufficient slots"), + } + } + + #[test] + fn test_prepare_rqd_run_frame_with_slots_required() { + // RunFrame message should include the slots_required value from the VirtualProc + let frame = create_test_dispatch_frame(); + let virtual_proc = VirtualProc { + proc_id: Uuid::new_v4(), + host_id: Uuid::new_v4(), + show_id: Uuid::new_v4(), + layer_id: Uuid::new_v4(), + job_id: Uuid::new_v4(), + frame_id: Uuid::new_v4(), + alloc_id: Uuid::new_v4(), + cores_reserved: CoreSize(1).with_multiplier(), + memory_reserved: ByteSize::gib(2), + gpus_reserved: 0, + gpu_memory_reserved: ByteSize::gb(0), + os: "linux".to_string(), + is_local_dispatch: false, + frame, + host_name: "somehost".to_string(), + slots_required: 3, + }; + + let result = RqdDispatcherService::prepare_rqd_run_frame(&virtual_proc); + + assert!(result.is_ok()); + let run_frame = result.unwrap(); + assert_eq!(run_frame.slots_required, 3); + } + + #[tokio::test] + async fn test_consume_host_virtual_resources_sequential_slot_consumption() { + // Each successive slot-based dispatch reduces available slots and resources + let mut frame = create_test_dispatch_frame(); + frame.slots_required = 1; + frame.min_cores = CoreSize(1); + frame.min_memory = ByteSize::gib(2); + frame.threadable = false; + + let mut host = create_test_host(); + host.concurrent_slots_limit = Some(4); + + // First dispatch + let (vp1, host_after_first) = RqdDispatcherService::consume_host_virtual_resources( + &frame, + &host, + ByteSize::gib(1), + ) + .await + .expect("first slot dispatch should succeed"); + + assert_eq!(vp1.slots_required, 1); + assert_eq!(host_after_first.running_slots_count, 1); + assert!(host_after_first.idle_memory < host.idle_memory); + + // Second dispatch on the updated host + let (vp2, host_after_second) = RqdDispatcherService::consume_host_virtual_resources( + &frame, + &host_after_first, + ByteSize::gib(1), + ) + .await + .expect("second slot dispatch should succeed"); + + assert_eq!(vp2.slots_required, 1); + assert_eq!(host_after_second.running_slots_count, 2); + assert!(host_after_second.idle_memory < host_after_first.idle_memory); + } + + #[tokio::test] + async fn test_consume_host_virtual_resources_slots_zero_vs_nonzero() { + // Verify slots_required = 0 and != 0 produce VirtualProcs with the correct field + let frame_no_slots = create_test_dispatch_frame(); // slots_required = 0 + + let mut frame_with_slots = create_test_dispatch_frame(); + frame_with_slots.slots_required = 5; + + let host = create_test_host(); + + let (vp_no_slots, _) = RqdDispatcherService::consume_host_virtual_resources( + &frame_no_slots, + &host, + ByteSize::gib(1), + ) + .await + .expect("dispatch without slots should succeed"); + + let (vp_with_slots, _) = RqdDispatcherService::consume_host_virtual_resources( + &frame_with_slots, + &host, + ByteSize::gib(1), + ) + .await + .expect("dispatch with slots should succeed"); + + assert_eq!(vp_no_slots.slots_required, 0); + assert_eq!(vp_with_slots.slots_required, 5); + } } diff --git a/rust/crates/scheduler/src/pipeline/matcher.rs b/rust/crates/scheduler/src/pipeline/matcher.rs index 2b4ee7c81..a277651fb 100644 --- a/rust/crates/scheduler/src/pipeline/matcher.rs +++ b/rust/crates/scheduler/src/pipeline/matcher.rs @@ -269,27 +269,23 @@ impl MatchingService { layer.show_id ); - // Clone only the minimal data needed for the validation closure - // These are needed because the closure must have 'static lifetime for actor messaging - let layer_id = layer.id; - let show_id = layer.show_id; let cores_requested = layer.cores_min; let allocation_service = self.allocation_service.clone(); let os = layer.str_os.clone(); + // Get a matching candidate let host_candidate = self .host_service .send(CheckOut { facility_id: layer.facility_id, show_id: layer.show_id, tags, - cores: cores_requested, - memory: layer.mem_min, + resource_request: layer.resource_request(), validation: move |host| { Self::validate_match( host, - &layer_id, - &show_id, + &layer.id, + &layer.show_id, cores_requested, &allocation_service, os.as_deref(),