Use SetProcessAffinityMask instead of SetProcessDefaultCpuSets (dotne…

…t#684)
caaavik-msft · Jan 12, 2024 · 73dda7b · 73dda7b
1 parent c5d29fa
commit 73dda7b
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 38 deletions.
diff --git a/src/Microsoft.Crank.Agent/Startup.cs b/src/Microsoft.Crank.Agent/Startup.cs
@@ -4603,31 +4603,6 @@ private static async Task<Process> StartProcess(string hostname, string benchmar
 
             var useWindowsLimiter = OperatingSystem == OperatingSystem.Windows && (job.MemoryLimitInBytes > 0 || job.CpuLimitRatio > 0 || !String.IsNullOrWhiteSpace(job.CpuSet));
 
-            // If useWindowsLimiter is true we need to use a proxy app that will get a JobObject attached such that the actual application can account for the limits
-            // when it starts. Otherwise the dotnet app would start while the JobOject is not yet defined.
-
-            if (useWindowsLimiter)
-            {
-                // This won't be necessary once https://github.com/dotnet/runtime/issues/94127 is implemented.
-                // At that point the process can be started suspended while the JobObject is assigned to it.
-
-                Console.WriteLine("Invoking JobObject Wrapper");
-
-                var filename = process.StartInfo.FileName;
-                var arguments = process.StartInfo.Arguments;
-
-                // The executable should be in the same folder as the agent since it references the Console project
-                // Use 'dotnet exec .dll' to use the current default dotnet version or the tests could fail if the .exe doesn't match what version is available locally
-                process.StartInfo.FileName = "dotnet";
-                process.StartInfo.Arguments = "exec" + " " + Path.Combine(Path.GetDirectoryName(Assembly.GetExecutingAssembly().Location), "Microsoft.Crank.JobObjectWrapper.dll") + " " + filename + " " + arguments;
-
-                // .NET doesn't respect a cpu affinity if a ratio is not set too. https://github.com/dotnet/runtime/issues/94364
-                if (!String.IsNullOrWhiteSpace(job.CpuSet) && job.CpuLimitRatio == 0)
-                {
-                    process.StartInfo.EnvironmentVariables.Add("DOTNET_PROCESSOR_COUNT", CalculateCpuList(job.CpuSet).Count.ToString(CultureInfo.InvariantCulture));
-                }
-            }
-
             stopwatch.Start();
             process.Start();
 
@@ -4644,7 +4619,7 @@ private static async Task<Process> StartProcess(string hostname, string benchmar
 
             if (useWindowsLimiter)
             {
-                var limiter = new WindowsLimiter(job, (uint)process.Id);
+                var limiter = new WindowsLimiter(job, process);
 
                 limiter.LimitProcess();
 

diff --git a/src/Microsoft.Crank.Agent/WindowsLimiter.cs b/src/Microsoft.Crank.Agent/WindowsLimiter.cs
@@ -2,12 +2,14 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
-using static Vanara.PInvoke.Kernel32;
-using Vanara.PInvoke;
-using Microsoft.Crank.Models;
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Linq;
+using System.Runtime.InteropServices;
+using Microsoft.Crank.Models;
+using Vanara.PInvoke;
+using static Vanara.PInvoke.Kernel32;
 
 namespace Microsoft.Crank.Agent
 {
@@ -19,13 +21,14 @@ public class WindowsLimiter : IDisposable
         private readonly SafeHJOB _safeJob;
 
         private readonly Job _job;
+        private readonly Process _process;
         private readonly List<int> _cpuSet;
 
-        public WindowsLimiter(Job job, uint processId) 
+        public WindowsLimiter(Job job, Process process) 
         {
             _job = job;
-
-            _safeProcess = OpenProcess(ACCESS_MASK.MAXIMUM_ALLOWED, false, processId);
+            _process = process;
+            _safeProcess = OpenProcess(ACCESS_MASK.MAXIMUM_ALLOWED, false, (uint)process.Id);
             _safeJob = CreateJobObject(null, $"{job.RunId}-{job.Service}");
 
             if (!String.IsNullOrWhiteSpace(job.CpuSet))
@@ -88,15 +91,21 @@ public void LimitProcess()
             if (_cpuSet != null && _cpuSet.Any())
             {
                 var ssi = GetSystemCpuSetInformation(_safeProcess).ToArray();
-                var cpuSets = _cpuSet.Select(i => ssi[i].CpuSet.Id).ToArray();
 
-                var result = SetProcessDefaultCpuSets(_safeProcess, cpuSets, (uint)cpuSets.Length);
-                Log.Info($"Limiting CpuSet ids: {String.Join(',', cpuSets)}, for ({_job.Service}:{_job.Id}): {(result ? "SUCCESS" : "FAILED")}");
+                Log.Info($"Limiting cpus: {String.Join(',', _cpuSet)}, for ({_job.Service}:{_job.Id}) Process: {_process.Id}");
 
-                foreach (var csi in ssi)
+                foreach (var c in _cpuSet)
                 {
+                    var csi = ssi[c];
                     Log.Info($"Id: {csi.CpuSet.Id}; NumaNodeIndex: {csi.CpuSet.NumaNodeIndex}; LogicalProcessorIndex: {csi.CpuSet.LogicalProcessorIndex}; CoreIndex: {csi.CpuSet.CoreIndex}; Group: {csi.CpuSet.Group}");
                 }
+
+                // Only supported on Linux and windows
+                if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows) || RuntimeInformation.IsOSPlatform(OSPlatform.Linux))
+                {
+                    var mask = (int)_cpuSet.Sum(x => Math.Pow(2, x));
+                    _process.ProcessorAffinity = new IntPtr(mask);
+                }
             }
 
             if (_hasJobObj)

diff --git a/src/Microsoft.Crank.JobOjectWrapper/Program.cs b/src/Microsoft.Crank.JobOjectWrapper/Program.cs
@@ -10,10 +10,11 @@
     Environment.Exit(-1);
 }
 
+var time = 1000;
 Console.WriteLine("Job Object Wrapper");
-Console.WriteLine("Waiting for Job Object to be setup...");
+Console.WriteLine($"Waiting for Job Object to be setup... {time}");
 
-await Task.Delay(1000);
+await Task.Delay(time);
 
 var process = new Process()
 {