schedulers/aws_batch: fix thread local sessions + raise error on missing memory resource

Author: d4l3k

torchx/schedulers/aws_batch_scheduler.py

@@ -38,7 +38,7 @@
 import threading
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Dict, Iterable, Mapping, Optional, Any, TYPE_CHECKING, Tuple
+from typing import Dict, Iterable, Mapping, Optional, Any, TYPE_CHECKING, Tuple, TypeVar, Callable
 
 import torchx
 import yaml
@@ -86,8 +86,8 @@ def _role_to_node_properties(idx: int, role: Role) -> Dict[str, object]:
     reqs.append({"type": "VCPU", "value": str(cpu)})
 
     mem = resource.memMB
-    if mem <= 0:
-        mem = 1000
+    if mem < 0:
+        raise ValueError(f"AWSBatchScheduler requires memMB to be set to a positive value, got {mem}")
     reqs.append({"type": "MEMORY", "value": str(mem)})
 
     if resource.gpu > 0:
@@ -157,17 +157,27 @@ def __repr__(self) -> str:
         return str(self)
 
 
-def _thread_local_session() -> "boto3.session.Session":
-    KEY = "torchx_aws_batch_session"
-    local = threading.local()
-    if hasattr(local, KEY):
-        # pyre-ignore[16]
-        return getattr(local, KEY)
+T = TypeVar('T')
+
+def _thread_local_cache(f: Callable[[], T]) -> Callable[[], T]:
+    local: threading.local = threading.local()
+    key: str = "value"
+    def wrapper() -> T:
+        if key in local.__dict__:
+            return local.__dict__[key]
+
+        v = f()
+        local.__dict__[key] = v
+        return v
+
+    return wrapper
+
+
+@_thread_local_cache
+def _local_session() -> "boto3.session.Session":
     import boto3.session
 
-    session = boto3.session.Session()
-    setattr(local, KEY, session)
-    return session
+    return boto3.session.Session()
 
 
 class AWSBatchScheduler(Scheduler, DockerWorkspace):
@@ -239,14 +249,14 @@ def __init__(
     def _client(self) -> Any:
         if self.__client:
             return self.__client
-        return _thread_local_session().client("batch")
+        return _local_session().client("batch")
 
     @property
     # pyre-fixme[3]: Return annotation cannot be `Any`.
     def _log_client(self) -> Any:
         if self.__log_client:
             return self.__log_client
-        return _thread_local_session().client("logs")
+        return _local_session().client("logs")
 
     def schedule(self, dryrun_info: AppDryRunInfo[BatchJob]) -> str:
         cfg = dryrun_info._cfg

torchx/schedulers/test/aws_batch_scheduler_test.py

@@ -8,13 +8,15 @@
 from contextlib import contextmanager
 from typing import Generator
 from unittest.mock import MagicMock, patch
+import threading
 
 import torchx
 from torchx import specs
 from torchx.schedulers.aws_batch_scheduler import (
     create_scheduler,
     AWSBatchScheduler,
     _role_to_node_properties,
+    _local_session,
 )
 
 
@@ -192,6 +194,11 @@ def test_volume_mounts(self) -> None:
             mounts=[
                 specs.VolumeMount(src="efsid", dst_path="/dst", read_only=True),
             ],
+            resource=specs.Resource(
+                cpu=1,
+                memMB=1000,
+                gpu=0,
+            ),
         )
         props = _role_to_node_properties(0, role)
         self.assertEqual(
@@ -396,3 +403,16 @@ def test_log_iter(self) -> None:
                 "foobar",
             ],
         )
+
+    def test_local_session(self) -> None:
+        a: object = _local_session()
+        self.assertIs(a, _local_session())
+
+        def worker() -> None:
+            b = _local_session()
+            self.assertIs(b, _local_session())
+            self.assertIsNot(a, b)
+
+        t = threading.Thread(target=worker)
+        t.start()
+        t.join()