code: def _test_stage_scheduling(
      self,
      cpus_per_worker,
      gpus_per_worker,
      num_tasks,
      resource_profile,
      expected_max_concurrent_tasks,
  ):
      conf = SparkConf()
      conf.setMaster(f"local-cluster[1,{cpus_per_worker},1024]").set(
          "spark.task.maxFailures", "1"
      )
  
      if gpus_per_worker:
          worker_res_config_file = os.path.join(
              self.temp_dir, "worker_res.json"
          )
          worker_res = [
              {
                  "id": {
                      "componentName": "spark.worker",
                      "resourceName": "gpu",
                  },
                  "addresses": [str(i) for i in range(gpus_per_worker)],
              }
          ]
          with open(worker_res_config_file, "w") as fp:
              json.dump(worker_res, fp)
  
          conf.set(
              "spark.worker.resource.gpu.amount", str(gpus_per_worker)
          )
          conf.set("spark.worker.resourcesFile", worker_res_config_file)
          conf.set(
              "spark.executor.resource.gpu.amount", str(gpus_per_worker)
          )
  
      self.sc = SparkContext(conf=conf)
      pids_output_dir = os.path.join(self.temp_dir, "pids")
      os.mkdir(pids_output_dir)
  
      def mapper(_):
          task_id = TaskContext.get().partitionId()
          pid_file_path = os.path.join(pids_output_dir, str(task_id))
          with open(pid_file_path, mode="w"):
              pass
          time.sleep(0.1)
          num_concurrent_tasks = len(os.listdir(pids_output_dir))
          time.sleep(1)
          os.remove(pid_file_path)
          return num_concurrent_tasks
  
      results = (
          self.sc.parallelize(range(num_tasks), num_tasks)
          .withResources(resource_profile)
          .map(mapper)
          .collect()
      )
      self.assertEqual(max(results), expected_max_concurrent_tasks)
  
file_path: spark/python/pyspark/tests/test_stage_sched.py
imports: import os
  import tempfile
  import unittest
  import time
  import shutil
  import json
  from pyspark import SparkConf, SparkContext
  from pyspark.resource.profile import ResourceProfileBuilder
  from pyspark.resource.requests import TaskResourceRequests
  from pyspark.taskcontext import TaskContext
label_desc: Use image node_id 3 for calling the StageSchedulingTest obj's underlying member method code with example usage: obj._test_stage_scheduling(cpus_per_worker, gpus_per_worker, num_tasks, resource_profile, expected_max_concurrent_tasks) and returns: num_concurrent_tasks