(torchx/schedulers) Remove redundantly setting dryruninfo._app and _cfg in k8s, gcp_batch, aws_batch, and docker scheduler subclasses since the scheduler abstract class sets it in the dryrun() API (#674)

Author: kiukchung

torchx/schedulers/aws_batch_scheduler.py

@@ -442,11 +442,7 @@ def _submit_dryrun(self, app: AppDef, cfg: AWSBatchOpts) -> AppDryRunInfo[BatchJ
             job_def=job_def,
             images_to_push=images_to_push,
         )
-        info = AppDryRunInfo(req, repr)
-        info._app = app
-        # pyre-fixme: AppDryRunInfo
-        info._cfg = cfg
-        return info
+        return AppDryRunInfo(req, repr)
 
     def _validate(self, app: AppDef, scheduler: str) -> None:
         # Skip validation step

torchx/schedulers/docker_scheduler.py

@@ -302,11 +302,7 @@ def _submit_dryrun(self, app: AppDef, cfg: DockerOpts) -> AppDryRunInfo[DockerJo
                     ]
                 req.containers.append(c)
 
-        info = AppDryRunInfo(req, repr)
-        info._app = app
-        # pyre-fixme: AppDryRunInfo
-        info._cfg = cfg
-        return info
+        return AppDryRunInfo(req, repr)
 
     def _validate(self, app: AppDef, scheduler: str) -> None:
         # Skip validation step

torchx/schedulers/gcp_batch_scheduler.py

@@ -313,11 +313,7 @@ def _submit_dryrun(
             job_def=job,
         )
 
-        info = AppDryRunInfo(req, repr)
-        info._app = app
-        # pyre-fixme: AppDryRunInfo
-        info._cfg = cfg
-        return info
+        return AppDryRunInfo(req, repr)
 
     def run_opts(self) -> runopts:
         opts = runopts()

torchx/schedulers/kubernetes_scheduler.py

@@ -615,11 +615,7 @@ def _submit_dryrun(
             resource=resource,
             images_to_push=images_to_push,
         )
-        info = AppDryRunInfo(req, repr)
-        info._app = app
-        # pyre-fixme: AppDryRunInfo
-        info._cfg = cfg
-        return info
+        return AppDryRunInfo(req, repr)
 
     def _validate(self, app: AppDef, scheduler: str) -> None:
         # Skip validation step

torchx/schedulers/test/gcp_batch_scheduler_test.py

@@ -69,7 +69,7 @@ def test_submit_dryrun(self) -> None:
         proj = "test-proj"
         loc = "us-west-1"
         cfg = GCPBatchOpts(project=proj, location=loc)
-        info = scheduler._submit_dryrun(app, cfg)
+        info = scheduler.submit_dryrun(app, cfg)
 
         req = info.request
         self.assertEqual(req.project, proj)
@@ -149,7 +149,7 @@ def test_submit_dryrun_throws(self) -> None:
         app.roles[0].resource.gpu = 3
         cfg = GCPBatchOpts(project="test-proj", location="us-west-1")
         with self.assertRaises(ValueError):
-            scheduler._submit_dryrun(app, cfg)
+            scheduler.submit_dryrun(app, cfg)
 
     def test_app_id_to_job_full_name(self) -> None:
         scheduler = create_scheduler("test")
@@ -369,7 +369,7 @@ def test_submit(self) -> None:
         # pyre-fixme: GCPBatchOpts type passed to resolve
         resolved_cfg = scheduler.run_opts().resolve(cfg)
         # pyre-fixme: _submit_dryrun expects GCPBatchOpts
-        info = scheduler._submit_dryrun(app, resolved_cfg)
+        info = scheduler.submit_dryrun(app, resolved_cfg)
         id = scheduler.schedule(info)
         self.assertEqual(id, "test-proj:us-central1:app-name-42")
         self.assertEqual(scheduler._client.create_job.call_count, 1)

torchx/schedulers/test/kubernetes_scheduler_test.py

@@ -213,7 +213,7 @@ def test_submit_dryrun(self) -> None:
             "torchx.schedulers.kubernetes_scheduler.make_unique"
         ) as make_unique_ctx:
             make_unique_ctx.return_value = "app-name-42"
-            info = scheduler._submit_dryrun(app, cfg)
+            info = scheduler.submit_dryrun(app, cfg)
 
         resource = str(info.request)
 
@@ -460,9 +460,8 @@ def test_rank0_env(self) -> None:
             "torchx.schedulers.kubernetes_scheduler.make_unique"
         ) as make_unique_ctx:
             make_unique_ctx.return_value = "app-name-42"
-            info = scheduler._submit_dryrun(app, cfg)
+            info = scheduler.submit_dryrun(app, cfg)
 
-        # pyre-fixme[16]; `object` has no attribute `__getitem__`.
         tasks = info.request.resource["spec"]["tasks"]
         container0 = tasks[0]["template"].spec.containers[0]
         self.assertIn("TORCHX_RANK0_HOST", container0.command)
@@ -486,7 +485,7 @@ def test_submit_dryrun_patch(self) -> None:
             "torchx.schedulers.kubernetes_scheduler.make_unique"
         ) as make_unique_ctx:
             make_unique_ctx.return_value = "app-name-42"
-            info = scheduler._submit_dryrun(app, cfg)
+            info = scheduler.submit_dryrun(app, cfg)
 
         self.assertIn("example.com/some/repo:testhash", str(info.request.resource))
         self.assertEqual(
@@ -509,11 +508,11 @@ def test_submit_dryrun_service_account(self) -> None:
                 "service_account": "srvacc",
             }
         )
-        info = scheduler._submit_dryrun(app, cfg)
+        info = scheduler.submit_dryrun(app, cfg)
         self.assertIn("'service_account_name': 'srvacc'", str(info.request.resource))
 
         del cfg["service_account"]
-        info = scheduler._submit_dryrun(app, cfg)
+        info = scheduler.submit_dryrun(app, cfg)
         self.assertIn("service_account_name': None", str(info.request.resource))
 
     def test_submit_dryrun_priority_class(self) -> None:
@@ -527,11 +526,11 @@ def test_submit_dryrun_priority_class(self) -> None:
             }
         )
 
-        info = scheduler._submit_dryrun(app, cfg)
+        info = scheduler.submit_dryrun(app, cfg)
         self.assertIn("'priorityClassName': 'high'", str(info.request.resource))
 
         del cfg["priority_class"]
-        info = scheduler._submit_dryrun(app, cfg)
+        info = scheduler.submit_dryrun(app, cfg)
         self.assertNotIn("'priorityClassName'", str(info.request.resource))
 
     @patch("kubernetes.client.CustomObjectsApi.create_namespaced_custom_object")
@@ -548,7 +547,7 @@ def test_submit(self, create_namespaced_custom_object: MagicMock) -> None:
             }
         )
 
-        info = scheduler._submit_dryrun(app, cfg)
+        info = scheduler.submit_dryrun(app, cfg)
         id = scheduler.schedule(info)
         self.assertEqual(id, "testnamespace:testid")
         call = create_namespaced_custom_object.call_args
@@ -577,7 +576,7 @@ def test_submit_job_name_conflict(
                 "queue": "testqueue",
             }
         )
-        info = scheduler._submit_dryrun(app, cfg)
+        info = scheduler.submit_dryrun(app, cfg)
         with self.assertRaises(ValueError):
             scheduler.schedule(info)
 
@@ -895,4 +894,4 @@ def test_dryrun(self) -> None:
         )
 
         with self.assertRaises(ModuleNotFoundError):
-            scheduler._submit_dryrun(app, cfg)
+            scheduler.submit_dryrun(app, cfg)

torchx/specs/api.py

@@ -657,7 +657,7 @@ def __init__(self, request: T, fmt: Callable[[T], str]) -> None:
         # Scheduler or Session implementations
         # and are back references to the parameters
         # to dryrun() that returned this AppDryRunInfo object
-        # thus they are set in Session.dryrun() and Scheduler.submit_dryrun()
+        # thus they are set in Runner.dryrun() and Scheduler.submit_dryrun()
         # manually rather than through constructor arguments
         # DO NOT create getters or make these public
         # unless there is a good reason to