From f31d76cb5ea2dc7eaf4e86bd5a3399a3ad7fe9f5 Mon Sep 17 00:00:00 2001 From: Aliaksandr Ivanou Date: Wed, 23 Jun 2021 17:33:47 -0700 Subject: [PATCH] Make torchx compatible with python 3.7 (#89) Summary: Pull Request resolved: https://github.com/pytorch/torchx/pull/89 Make torchx compatible with python 3.7 Differential Revision: D29348727 fbshipit-source-id: 5351e47a7783d9fb09124f7d712b2d714efe1098 --- torchx/schedulers/slurm_scheduler.py | 17 +++++++++-------- torchx/schedulers/test/slurm_scheduler_test.py | 7 +++---- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/torchx/schedulers/slurm_scheduler.py b/torchx/schedulers/slurm_scheduler.py index cb53d84fd..9663757bb 100644 --- a/torchx/schedulers/slurm_scheduler.py +++ b/torchx/schedulers/slurm_scheduler.py @@ -68,14 +68,15 @@ class SlurmReplicaRequest: @classmethod def from_role(cls, role: Role, cfg: RunConfig) -> "SlurmReplicaRequest": opts = {k: str(v) for k, v in cfg.cfgs.items()} - - if (resource := role.resource) != NONE: - if (cpu := resource.cpu) > 0: - opts["cpus-per-task"] = str(cpu) - if (memMB := resource.memMB) > 0: - opts["mem"] = str(memMB) - if (gpu := resource.gpu) > 0: - opts["gpus-per-task"] = str(gpu) + resource = role.resource + + if resource != NONE: + if resource.cpu > 0: + opts["cpus-per-task"] = str(resource.cpu) + if resource.memMB > 0: + opts["mem"] = str(resource.memMB) + if resource.gpu > 0: + opts["gpus-per-task"] = str(resource.gpu) return cls( dir=role.image, diff --git a/torchx/schedulers/test/slurm_scheduler_test.py b/torchx/schedulers/test/slurm_scheduler_test.py index 61ce27f70..3d2401c93 100644 --- a/torchx/schedulers/test/slurm_scheduler_test.py +++ b/torchx/schedulers/test/slurm_scheduler_test.py @@ -138,10 +138,9 @@ def test_run_multi_role(self, run: MagicMock) -> None: self.assertEqual(app_id, "1234") self.assertEqual(run.call_count, 1) - self.assertEqual( - run.call_args.kwargs, {"stdout": subprocess.PIPE, "check": True} - ) - (args,) = run.call_args.args + args, kwargs = run.call_args + self.assertEqual(kwargs, {"stdout": subprocess.PIPE, "check": True}) + (args,) = args self.assertEqual(len(args), 9) self.assertEqual(args[:4], ["sbatch", "--parsable", "--job-name", "foo"]) self.assertTrue(args[4].endswith("role-0-a-0.sh"))