[rllib/tune] Add test for fractional gpu support in xray mode; add rllib support for fractional gpu (#2768)

* frac gpu

* doc

* Update rllib-training.rst

* yapf

* remove xray

python/ray/rllib/agents/a3c/a2c.py

@@ -34,6 +34,6 @@ class A2CAgent(A3CAgent):
         cf = merge_dicts(cls._default_config, config)
         return Resources(
             cpu=1,
-            gpu=1 if cf["gpu"] else 0,
+            gpu=cf["gpu_fraction"] if cf["gpu"] else 0,
             extra_cpu=cf["num_workers"],
             extra_gpu=cf["use_gpu_for_workers"] and cf["num_workers"] or 0)

python/ray/rllib/agents/agent.py

@@ -64,6 +64,8 @@ COMMON_CONFIG = {
     "compress_observations": False,
     # Whether to write episode stats and videos to the agent log dir
     "monitor": False,
+    # Allocate a fraction of a GPU instead of one (e.g., 0.3 GPUs)
+    "gpu_fraction": 1,
 
     # === Multiagent ===
     "multiagent": {

python/ray/rllib/agents/bc/bc.py

@@ -53,12 +53,12 @@ class BCAgent(Agent):
     def default_resource_request(cls, config):
         cf = merge_dicts(cls._default_config, config)
         if cf["use_gpu_for_workers"]:
-            num_gpus_per_worker = 1
+            num_gpus_per_worker = cf["gpu_fraction"]
         else:
             num_gpus_per_worker = 0
         return Resources(
             cpu=1,
-            gpu=cf["gpu"] and 1 or 0,
+            gpu=cf["gpu"] and cf["gpu_fraction"] or 0,
             extra_cpu=cf["num_workers"],
             extra_gpu=num_gpus_per_worker * cf["num_workers"])
 

python/ray/rllib/agents/ddpg/apex.py

@@ -48,7 +48,7 @@ class ApexDDPGAgent(DDPGAgent):
         cf = merge_dicts(cls._default_config, config)
         return Resources(
             cpu=1 + cf["optimizer"]["num_replay_buffer_shards"],
-            gpu=cf["gpu"] and 1 or 0,
+            gpu=cf["gpu"] and cf["gpu_fraction"] or 0,
             extra_cpu=cf["num_cpus_per_worker"] * cf["num_workers"],
             extra_gpu=cf["num_gpus_per_worker"] * cf["num_workers"])
 

python/ray/rllib/agents/dqn/apex.py

@@ -47,7 +47,7 @@ class ApexAgent(DQNAgent):
         cf = merge_dicts(cls._default_config, config)
         return Resources(
             cpu=1 + cf["optimizer"]["num_replay_buffer_shards"],
-            gpu=cf["gpu"] and 1 or 0,
+            gpu=cf["gpu"] and cf["gpu_fraction"] or 0,
             extra_cpu=cf["num_cpus_per_worker"] * cf["num_workers"],
             extra_gpu=cf["num_gpus_per_worker"] * cf["num_workers"])
 

python/ray/rllib/agents/dqn/dqn.py

@@ -132,7 +132,7 @@ class DQNAgent(Agent):
         cf = merge_dicts(cls._default_config, config)
         return Resources(
             cpu=1,
-            gpu=cf["gpu"] and 1 or 0,
+            gpu=cf["gpu"] and cf["gpu_fraction"] or 0,
             extra_cpu=cf["num_cpus_per_worker"] * cf["num_workers"],
             extra_gpu=cf["num_gpus_per_worker"] * cf["num_workers"])
 

python/ray/rllib/agents/impala/impala.py

@@ -69,7 +69,7 @@ class ImpalaAgent(Agent):
         cf = dict(cls._default_config, **config)
         return Resources(
             cpu=1,
-            gpu=cf["gpu"] and 1 or 0,
+            gpu=cf["gpu"] and cf["gpu_fraction"] or 0,
             extra_cpu=cf["num_cpus_per_worker"] * cf["num_workers"],
             extra_gpu=cf["num_gpus_per_worker"] * cf["num_workers"])