diff --git a/mantis-runtime-loader/src/main/java/io/mantisrx/runtime/loader/config/WorkerConfiguration.java b/mantis-runtime-loader/src/main/java/io/mantisrx/runtime/loader/config/WorkerConfiguration.java index 52f460006..788197015 100644 --- a/mantis-runtime-loader/src/main/java/io/mantisrx/runtime/loader/config/WorkerConfiguration.java +++ b/mantis-runtime-loader/src/main/java/io/mantisrx/runtime/loader/config/WorkerConfiguration.java @@ -155,6 +155,18 @@ default Time getHeartbeatInterval() { @DefaultNull File getLocalStorageDir(); + @Config("mantis.taskexecutor.hardware.cpu-cores") + @DefaultNull + Double getCpuCores(); + + @Config("mantis.taskexecutor.hardware.memory-in-mb") + @DefaultNull + Double getMemoryInMB(); + + @Config("mantis.taskexecutor.hardware.disk-in-mb") + @DefaultNull + Double getDiskInMB(); + @Config("mantis.taskexecutor.hardware.network-bandwidth-in-mb") @Default(value = "128.0") double getNetworkBandwidthInMB(); diff --git a/mantis-runtime-loader/src/main/java/io/mantisrx/runtime/loader/config/WorkerConfigurationUtils.java b/mantis-runtime-loader/src/main/java/io/mantisrx/runtime/loader/config/WorkerConfigurationUtils.java index 1a50a1157..bcc5ab7a0 100644 --- a/mantis-runtime-loader/src/main/java/io/mantisrx/runtime/loader/config/WorkerConfigurationUtils.java +++ b/mantis-runtime-loader/src/main/java/io/mantisrx/runtime/loader/config/WorkerConfigurationUtils.java @@ -53,6 +53,9 @@ public static WorkerConfigurationWritable toWrit .metricsPort(configSource.getMetricsPort()) .metricsPublisher(configSource.getMetricsPublisher()) .metricsPublisherFrequencyInSeconds(configSource.getMetricsPublisherFrequencyInSeconds()) + .cpuCores(configSource.getCpuCores()) + .memoryInMB(configSource.getMemoryInMB()) + .diskInMB(configSource.getDiskInMB()) .networkBandwidthInMB(configSource.getNetworkBandwidthInMB()) .sinkPort(configSource.getSinkPort()) .taskExecutorId(configSource.getTaskExecutorId()) diff --git a/mantis-runtime-loader/src/main/java/io/mantisrx/runtime/loader/config/WorkerConfigurationWritable.java b/mantis-runtime-loader/src/main/java/io/mantisrx/runtime/loader/config/WorkerConfigurationWritable.java index 37ed708a2..6c885c88d 100644 --- a/mantis-runtime-loader/src/main/java/io/mantisrx/runtime/loader/config/WorkerConfigurationWritable.java +++ b/mantis-runtime-loader/src/main/java/io/mantisrx/runtime/loader/config/WorkerConfigurationWritable.java @@ -65,6 +65,9 @@ public class WorkerConfigurationWritable implements WorkerConfiguration { URI blobStoreArtifactDir; File localStorageDir; + Double cpuCores; + Double memoryInMB; + Double diskInMB; double networkBandwidthInMB; String taskExecutorAttributesStr; int asyncHttpClientMaxConnectionsPerHost; @@ -243,6 +246,21 @@ public File getLocalStorageDir() { return this.localStorageDir; } + @Override + public Double getCpuCores() { + return this.cpuCores; + } + + @Override + public Double getMemoryInMB() { + return this.memoryInMB; + } + + @Override + public Double getDiskInMB() { + return this.diskInMB; + } + @Override public double getNetworkBandwidthInMB() { return this.networkBandwidthInMB; diff --git a/mantis-server/mantis-server-agent/src/main/java/io/mantisrx/server/agent/MachineDefinitionUtils.java b/mantis-server/mantis-server-agent/src/main/java/io/mantisrx/server/agent/MachineDefinitionUtils.java index 845997e3d..54f737cf8 100644 --- a/mantis-server/mantis-server-agent/src/main/java/io/mantisrx/server/agent/MachineDefinitionUtils.java +++ b/mantis-server/mantis-server-agent/src/main/java/io/mantisrx/server/agent/MachineDefinitionUtils.java @@ -17,8 +17,17 @@ import io.mantisrx.common.WorkerPorts; import io.mantisrx.runtime.MachineDefinition; +import io.mantisrx.runtime.loader.config.WorkerConfiguration; +import java.util.Optional; public class MachineDefinitionUtils { + /** + * Creates a MachineDefinition object using system parameters. + * + * @param workerPorts Worker ports object. + * @param networkBandwidthInMB Network bandwidth capacity in MB. + * @return MachineDefinition object. + */ public static MachineDefinition sys(WorkerPorts workerPorts, double networkBandwidthInMB) { return new MachineDefinition( Hardware.getNumberCPUCores(), @@ -27,4 +36,34 @@ public static MachineDefinition sys(WorkerPorts workerPorts, double networkBandw Hardware.getSizeOfDisk() / 1024.0 / 1024.0, workerPorts.getNumberOfPorts()); } + + /** + * Creates a MachineDefinition object from the worker configuration, falling back to system parameters. + * + * @param workerConfiguration Worker configuration object. + * @param workerPorts Worker ports object. + * @return MachineDefinition object. + */ + public static MachineDefinition from(WorkerConfiguration workerConfiguration, WorkerPorts workerPorts) { + return fromWorkerConfiguration(workerConfiguration, workerPorts).orElseGet(() -> sys(workerPorts, workerConfiguration.getNetworkBandwidthInMB())); + } + + /** + * Attempts to create a MachineDefinition from the WorkerConfiguration if all necessary parameters are present. + * + * @param workerConfiguration WorkerConfiguration object. + * @param workerPorts Worker ports object. + * @return An Optional of MachineDefinition if all parameters are present in WorkerConfiguration, else empty Optional. + */ + private static Optional fromWorkerConfiguration(WorkerConfiguration workerConfiguration, WorkerPorts workerPorts) { + if (workerConfiguration.getCpuCores() != null && workerConfiguration.getMemoryInMB() != null && workerConfiguration.getDiskInMB() != null) { + return Optional.of(new MachineDefinition( + workerConfiguration.getCpuCores(), + workerConfiguration.getMemoryInMB(), + workerConfiguration.getNetworkBandwidthInMB(), + workerConfiguration.getDiskInMB(), + workerPorts.getNumberOfPorts())); + } + return Optional.empty(); + } } diff --git a/mantis-server/mantis-server-agent/src/main/java/io/mantisrx/server/agent/TaskExecutor.java b/mantis-server/mantis-server-agent/src/main/java/io/mantisrx/server/agent/TaskExecutor.java index be3ec4cf2..5fde439f6 100644 --- a/mantis-server/mantis-server-agent/src/main/java/io/mantisrx/server/agent/TaskExecutor.java +++ b/mantis-server/mantis-server-agent/src/main/java/io/mantisrx/server/agent/TaskExecutor.java @@ -157,8 +157,7 @@ public TaskExecutor( workerConfiguration.getDebugPort(), workerConfiguration.getConsolePort(), workerConfiguration.getCustomPort(), workerConfiguration.getSinkPort()); - MachineDefinition machineDefinition = - MachineDefinitionUtils.sys(workerPorts, workerConfiguration.getNetworkBandwidthInMB()); + MachineDefinition machineDefinition = MachineDefinitionUtils.from(workerConfiguration, workerPorts); String hostName = workerConfiguration.getExternalAddress(); this.taskExecutorRegistration =