SlurmGCP. Fix type annotations problems

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/load_bq.py

@@ -26,7 +26,7 @@
 import util
 from google.api_core import exceptions, retry
 from google.cloud import bigquery as bq
-from google.cloud.bigquery import SchemaField
+from google.cloud.bigquery import SchemaField # type: ignore
 from util import lookup, run
 
 SACCT = "sacct"
@@ -175,7 +175,8 @@ def schema_field(field_name, data_type, description, required=False):
 # creating the job rows
 job_schema = {field.name: field for field in schema_fields}
 # Order is important here, as that is how they are parsed from sacct output
-Job = namedtuple("Job", job_schema.keys())
+Job = namedtuple("Job", job_schema.keys()) # type: ignore 
+# ... see https://github.com/python/mypy/issues/848
 
 client = bq.Client(
     project=lookup().cfg.project,

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/resume.py

@@ -41,7 +41,7 @@
     trim_self_link,
     wait_for_operation,
 )
-from util import lookup, NSDict
+from util import lookup, NSDict, ReservationDetails
 import tpu
 
 log = logging.getLogger()
@@ -104,6 +104,7 @@ def instance_properties(nodeset:object, model:str, placement_group:Optional[str]
         update_reservation_props(reservation, props, placement_group)
 
     if (fr := lookup().future_reservation(nodeset)) and fr.specific:
+        assert fr.active_reservation
         update_reservation_props(fr.active_reservation, props, placement_group)
 
     if props.resourcePolicies:
@@ -119,7 +120,7 @@ def instance_properties(nodeset:object, model:str, placement_group:Optional[str]
     props.update(nodeset.get("instance_properties") or {})
     return props
 
-def update_reservation_props(reservation:object, props:object, placement_group:Optional[str]) -> None:
+def update_reservation_props(reservation:ReservationDetails, props:object, placement_group:Optional[str]) -> None:
     props.reservationAffinity = {
         "consumeReservationType": "SPECIFIC_RESERVATION",
         "key": f"compute.{util.universe_domain()}/reservation-name",

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/setup_network_storage.py

@@ -24,7 +24,7 @@
 import shutil
 from pathlib import Path
 from concurrent.futures import as_completed
-from addict import Dict as NSDict
+from addict import Dict as NSDict # type: ignore
 
 import util
 from util import lookup, run, dirs, separate

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/sort_nodes.py

@@ -53,12 +53,12 @@ def order(paths: List[List[str]]) -> List[str]:
     if not paths: return []
     class Vert:
         "Represents a vertex in a *network* tree."
-        def __init__(self, name: str, parent: "Vert"):
+        def __init__(self, name: str, parent: Optional["Vert"]):
             self.name = name
             self.parent = parent
             # Use `OrderedDict` to preserve insertion order
             # TODO: once we move to Python 3.7+ use regular `dict` since it has the same guarantee
-            self.children = OrderedDict()
+            self.children: OrderedDict = OrderedDict()
 
     # build a tree, children are ordered by insertion order
     root = Vert("", None)
@@ -107,7 +107,7 @@ def to_hostnames(nodelist: str) -> List[str]:
     return [n.decode("utf-8") for n in out.splitlines()]
 
 
-def get_instances(node_names: List[str]) -> Dict[str, object]:
+def get_instances(node_names: List[str]) -> Dict[str, Optional[Instance]]:
     fmt = (
         "--format=csv[no-heading,separator=','](zone,resourceStatus.physicalHost,name)"
     )

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/tests/test_conf.py

@@ -16,7 +16,7 @@
 from mock import Mock
 from common import TstNodeset, TstCfg, TstMachineConf, TstTemplateInfo
 
-import addict
+import addict # type: ignore
 import conf
 import util
 

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/tests/test_topology.py

@@ -203,6 +203,6 @@ def test_gen_topology_conf_update():
         (["z/n-0", "z/n-1", "z/n-2", "z/n-3", "z/n-4", "z/n-10"], ['n-0', 'n-1', 'n-2', 'n-3', 'n-4', 'n-10']),
         (["y/n-0", "z/n-1", "x/n-2", "x/n-3", "y/n-4", "g/n-10"], ['n-0', 'n-4', 'n-1', 'n-2', 'n-3', 'n-10']),
     ])
-def test_sort_nodes_order(paths: list[list[str]], expected: list[str]) -> None:
-    paths = [l.split("/") for l in paths]
-    assert sort_nodes.order(paths) == expected
+def test_sort_nodes_order(paths: list[str], expected: list[str]) -> None:
+    paths_expanded = [l.split("/") for l in paths]
+    assert sort_nodes.order(paths_expanded) == expected

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/tests/test_util.py

@@ -318,7 +318,7 @@ def test_parse_job_info(job_info, expected_job):
 @pytest.mark.parametrize(
     "node,state,want",
     [
-        ("c-n-2", NodeState("DOWN", {}), NodeState("DOWN", {})), # happy scenario
+        ("c-n-2", NodeState("DOWN", frozenset([])), NodeState("DOWN", frozenset([]))), # happy scenario
         ("c-d-vodoo", None, None), # dynamic nodeset
         ("c-x-44", None, None), # unknown(removed) nodeset
         ("c-n-7", None, None), # Out of bounds: c-n-[0-4] - downsized nodeset
@@ -341,7 +341,8 @@ def test_node_state(node: str, state: Optional[NodeState], want: NodeState | Non
             "d": TstNodeset()},
     )
     lkp = util.Lookup(cfg)
-    lkp.slurm_nodes = lambda: {node: state} if state else {}
+    lkp.slurm_nodes = lambda: {node: state} if state else {} # type: ignore[assignment]
+    # ... see https://github.com/python/typeshed/issues/6347    
 
     if  type(want) is type and issubclass(want, Exception):
         with pytest.raises(want):