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

Author: d4l3k

torchx/schedulers/aws_batch_scheduler.py

@@ -38,7 +38,17 @@
 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 +96,10 @@ 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 +169,29 @@ 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 +263,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

@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import threading
 import unittest
 from contextlib import contextmanager
 from typing import Generator
@@ -15,6 +16,7 @@
     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()