From caae73d4136ec7e73b289ebfa767a959ebef248a Mon Sep 17 00:00:00 2001
From: Emilio Palumbo <emiliopalumbo@gmail.com>
Date: Tue, 11 Feb 2014 18:27:21 +0100
Subject: [PATCH] Add opttion to select memory unit for SGE

---
 jip/cluster.py       | 10 +++++++++-
 test/test_cluster.py | 25 +++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/jip/cluster.py b/jip/cluster.py
index 52fcafd..8f97d4c 100644
--- a/jip/cluster.py
+++ b/jip/cluster.py
@@ -382,6 +382,7 @@ def __init__(self):
         self.threads_pe = sge_cfg.get('threads_pe', None)
         self.mem_limit = sge_cfg.get('mem_limit', 'virtual_free')
         self.time_limit = sge_cfg.get('time_limit', 's_rt')
+        self.mem_unit = sge_cfg.get('mem_unit','M').upper()
 
     def resolve_log(self, job, path):
         if path is None:
@@ -444,7 +445,7 @@ def submit(self, job):
         if job.working_directory:
             cmd.extend(["-wd", job.working_directory])
         if job.max_memory > 0:
-            cmd.extend(["-l", '%s=%sM' % (self.mem_limit, str(job.max_memory))])
+            cmd.extend(["-l", '%s=%s' % (self.mem_limit, self._sge_mem(job.max_memory))])
         if job.account:
             cmd.extend(["-A", str(job.account)])
         if job.extra is not None:
@@ -489,6 +490,13 @@ def submit(self, job):
         match = re.search(expr, out)
         job.job_id = match.group('job_id')
 
+    def _sge_mem(self, mem):
+        if self.mem_unit == 'G':
+            mem = mem / 1024
+        if self.mem_unit == 'K':
+            mem = mem * 1024
+        return "%s%s" % (mem, self.mem_unit)
+
 
 class PBS(Cluster):
     """PBS/Torque extension of the Cluster implementation.
diff --git a/test/test_cluster.py b/test/test_cluster.py
index 333d33f..f718f25 100644
--- a/test/test_cluster.py
+++ b/test/test_cluster.py
@@ -45,3 +45,28 @@ def test_sge_threads_pe_loading():
     }
     sge = cl.SGE()
     assert sge.threads_pe == 'threads'
+
+@pytest.mark.parametrize("unit,op", [
+    ('G','/ 1024'),
+    ('g','/ 1024'),
+    ('K','* 1024'),
+    ('k','* 1024'),
+    ('M',''),
+    ('m','')
+])
+
+def test_sge_mem_unit_loading(unit,op):
+    mem = 32768
+    jip.config.config['sge'] = {
+        "mem_unit": unit
+    }
+    sge = cl.SGE()
+    assert sge.mem_unit == unit.upper()
+    assert sge._sge_mem(mem) == ('%s%s' % (eval("%s%s" % (mem,op)), unit.upper()))
+
+def test_sge_mem_unit_default():
+    mem = 32768
+    sge = cl.SGE()
+    assert sge.mem_unit == 'M'
+    assert sge._sge_mem(mem) == '32768M'
+