use TaskContext.get to avoid creating instance in unit tests

tensorflowonspark/TFSparkNode.py

@@ -189,20 +189,21 @@ def _get_gpus(cluster_spec=None):
       # note: num_gpus arg is only used (if supplied) to limit/truncate visible devices
       if _has_spark_resource_api():
         from pyspark import TaskContext
-        context = TaskContext()
-        resources = context.resources()
-        if resources and 'gpu' in resources:
-          # get all GPUs assigned by resource manager
-          gpus = context.resources()['gpu'].addresses
-          logger.info("Spark gpu resources: {}".format(gpus))
-          if user_requested:
-            if requested_gpus < len(gpus):
-              # override/truncate list, if explicitly configured
-              logger.warn("Requested {} GPU(s), but {} available".format(requested_gpus, len(gpus)))
-              gpus = gpus[:requested_gpus]
-          else:
-            # implicitly requested by Spark 3
-            requested_gpus = len(gpus)
+        context = TaskContext.get()
+        if context:
+          resources = context.resources()
+          if resources and 'gpu' in resources:
+            # get all GPUs assigned by resource manager
+            gpus = context.resources()['gpu'].addresses
+            logger.info("Spark gpu resources: {}".format(gpus))
+            if user_requested:
+              if requested_gpus < len(gpus):
+                # override/truncate list, if explicitly configured
+                logger.warn("Requested {} GPU(s), but {} available".format(requested_gpus, len(gpus)))
+                gpus = gpus[:requested_gpus]
+            else:
+              # implicitly requested by Spark 3
+              requested_gpus = len(gpus)
 
       # if not in K8s pod and GPUs available, just use original allocation code (defaulting to 1 GPU if available)
       # Note: for K8s, there is a bug with the Nvidia device_plugin which can show GPUs for non-GPU pods that are hosted on GPU nodes
@@ -348,15 +349,15 @@ def _get_gpus(cluster_spec=None):
         port = tmp_sock.getsockname()[1]
 
       node_meta = {
-          'executor_id': executor_id,
-          'host': host,
-          'job_name': job_name,
-          'task_index': task_index,
-          'port': port,
-          'tb_pid': tb_pid,
-          'tb_port': tb_port,
-          'addr': addr,
-          'authkey': authkey
+        'executor_id': executor_id,
+        'host': host,
+        'job_name': job_name,
+        'task_index': task_index,
+        'port': port,
+        'tb_pid': tb_pid,
+        'tb_port': tb_port,
+        'addr': addr,
+        'authkey': authkey
       }
       # register node metadata with server
       logger.info("TFSparkNode.reserve: {0}".format(node_meta))

tests/test_TFSparkNode.py

@@ -112,6 +112,7 @@ def test_gpu_spark_available(self, mock_available, mock_get_gpus, mock_spark_res
     mock_available.return_value = True
     mock_get_gpus.return_value = ['0']
     mock_spark_resources.return_value = True
+    mock_context.get.return_value = mock_context.return_value
     mock_context_instance = mock_context.return_value
     mock_context_instance.resources.return_value = {'gpu': type("ResourceInformation", (object,), {"addresses": ['0']})}