Allow using torchx_ env vars to set scheduler params

Summary: Scheduler params currently can only be set through the programmatic API and not through this is not useful for cases like scheduling on mast rc cluster. This diff now lets you do that.

Differential Revision: D57640022

Author: manav-a

torchx/runner/api.py

@@ -98,7 +98,12 @@ def __init__(
         """
         self._name: str = name
         self._scheduler_factories = scheduler_factories
-        self._scheduler_params: Dict[str, object] = scheduler_params or {}
+        self._scheduler_params: Dict[str, Any] = scheduler_params or {}
+        for key, value in os.environ.items():
+            lower_case_key = key.lower()
+            if lower_case_key.startswith("torchx_"):
+                self._scheduler_params[lower_case_key] = value
+
         # pyre-fixme[24]: SchedulerOpts is a generic, and we don't have access to the corresponding type
         self._scheduler_instances: Dict[str, Scheduler] = {}
         self._apps: Dict[AppHandle, AppDef] = {}

torchx/runner/test/api_test.py

@@ -10,14 +10,15 @@
 import datetime
 import os
 from contextlib import contextmanager
-from typing import Generator, List, Mapping, Optional
+from typing import cast, Generator, List, Mapping, Optional
 from unittest.mock import MagicMock, patch
 
 from torchx.runner import get_runner, Runner
+from torchx.schedulers import SchedulerFactory
 from torchx.schedulers.api import DescribeAppResponse, ListAppResponse, Scheduler
 from torchx.schedulers.local_scheduler import (
+    create_scheduler,
     LocalDirectoryImageProvider,
-    LocalScheduler,
 )
 from torchx.specs import AppDryRunInfo, CfgVal
 from torchx.specs.api import (
@@ -64,7 +65,7 @@ def setUp(self) -> None:
     def get_runner(self) -> Generator[Runner, None, None]:
         with Runner(
             SESSION_NAME,
-            scheduler_factories={"local_dir": LocalScheduler},
+            scheduler_factories={"local_dir": cast(SchedulerFactory, create_scheduler)},
             scheduler_params={
                 "image_provider_class": LocalDirectoryImageProvider,
             },
@@ -79,14 +80,14 @@ def test_validate_no_roles(self, _) -> None:
 
     def test_validate_no_resource(self, _) -> None:
         with self.get_runner() as runner:
+            role = Role(
+                "no resource",
+                image="no_image",
+                entrypoint="echo",
+                args=["hello_world"],
+            )
+            app = AppDef("no resource", roles=[role])
             with self.assertRaises(ValueError):
-                role = Role(
-                    "no resource",
-                    image="no_image",
-                    entrypoint="echo",
-                    args=["hello_world"],
-                )
-                app = AppDef("no resource", roles=[role])
                 runner.run(app, scheduler="local_dir")
 
     def test_validate_invalid_replicas(self, _) -> None:
@@ -129,7 +130,7 @@ def test_dryrun(self, _) -> None:
         }
         with Runner(
             name=SESSION_NAME,
-            scheduler_factories={"local_dir": lambda name: scheduler_mock},
+            scheduler_factories={"local_dir": lambda name, **kwargs: scheduler_mock},
         ) as runner:
             role = Role(
                 name="touch",
@@ -149,7 +150,7 @@ def test_dryrun_env_variables(self, _) -> None:
         scheduler_mock = MagicMock()
         with Runner(
             name=SESSION_NAME,
-            scheduler_factories={"local_dir": lambda name: scheduler_mock},
+            scheduler_factories={"local_dir": lambda name, **kwargs: scheduler_mock},
         ) as runner:
             role1 = Role(
                 name="echo1",
@@ -178,7 +179,7 @@ def test_dryrun_trackers_parent_run_id_as_paramenter(self, _) -> None:
         expected_parent_run_id = "123"
         with Runner(
             name=SESSION_NAME,
-            scheduler_factories={"local_dir": lambda name: scheduler_mock},
+            scheduler_factories={"local_dir": lambda name, **kwargs: scheduler_mock},
         ) as runner:
             role1 = Role(
                 name="echo1",
@@ -217,7 +218,7 @@ def test_dryrun_setup_trackers(self, config_trackers_mock: MagicMock, _) -> None
 
         with Runner(
             name=SESSION_NAME,
-            scheduler_factories={"local_dir": lambda name: scheduler_mock},
+            scheduler_factories={"local_dir": lambda name, **kwargs: scheduler_mock},
         ) as runner:
             role1 = Role(
                 name="echo1",
@@ -265,7 +266,7 @@ def test_dryrun_setup_trackers_as_env_variable(self, _) -> None:
 
         with Runner(
             name=SESSION_NAME,
-            scheduler_factories={"local_dir": lambda name: scheduler_mock},
+            scheduler_factories={"local_dir": lambda name, **kwargs: scheduler_mock},
         ) as runner:
             role1 = Role(
                 name="echo1",
@@ -371,7 +372,7 @@ def test_describe(self, _) -> None:
                 name="sleep",
                 image=str(self.tmpdir),
                 resource=resource.SMALL,
-                entrypoint="sleep.sh",
+                entrypoint="sleep",
                 args=["60"],
             )
             app = AppDef("sleeper", roles=[role])
@@ -387,7 +388,7 @@ def test_status(self, _) -> None:
                 name="sleep",
                 image=str(self.tmpdir),
                 resource=resource.SMALL,
-                entrypoint="sleep.sh",
+                entrypoint="sleep",
                 args=["60"],
             )
             app = AppDef("sleeper", roles=[role])
@@ -414,7 +415,7 @@ def test_status_ui_url(self, json_dumps_mock: MagicMock, _) -> None:
 
         with Runner(
             name="test_ui_url_session",
-            scheduler_factories={"local_dir": lambda name: mock_scheduler},
+            scheduler_factories={"local_dir": lambda name, **kwargs: mock_scheduler},
         ) as runner:
             role = Role(
                 "ignored",
@@ -438,7 +439,7 @@ def test_status_structured_msg(self, json_dumps_mock: MagicMock, _) -> None:
 
         with Runner(
             name="test_structured_msg",
-            scheduler_factories={"local_dir": lambda name: mock_scheduler},
+            scheduler_factories={"local_dir": lambda name, **kwargs: mock_scheduler},
         ) as runner:
             role = Role(
                 "ignored",
@@ -485,7 +486,7 @@ def test_log_lines(self, _) -> None:
 
         with Runner(
             name=SESSION_NAME,
-            scheduler_factories={"local_dir": lambda name: scheduler_mock},
+            scheduler_factories={"local_dir": lambda name, **kwargs: scheduler_mock},
         ) as runner:
             role_name = "trainer"
             replica_id = 2
@@ -605,7 +606,7 @@ def test_runner_context_manager_with_error(self, _) -> None:
             return_value={"local_dir": lambda name: mock_scheduler},
         ):
             with self.assertRaisesRegex(RuntimeError, "foobar"):
-                with get_runner() as runner:
+                with get_runner():
                     raise RuntimeError("foobar")
 
     def test_runner_manual_close(self, _) -> None:

torchx/schedulers/local_scheduler.py

@@ -1184,11 +1184,12 @@ def create_scheduler(
     session_name: str,
     cache_size: int = 100,
     extra_paths: Optional[List[str]] = None,
+    image_provider_class: Callable[[LocalOpts], ImageProvider] = CWDImageProvider,
     **kwargs: Any,
 ) -> LocalScheduler:
     return LocalScheduler(
         session_name=session_name,
-        image_provider_class=CWDImageProvider,
+        image_provider_class=image_provider_class,
         cache_size=cache_size,
         extra_paths=extra_paths,
     )