Fix typing issues

lxm3/contrib/wandb.py

@@ -19,6 +19,7 @@
     experiment.add(log_to_wandb(xm_cluster.ArrayJob(executable, ...)))
 
 """
+
 import copy
 import functools
 import logging
@@ -67,7 +68,7 @@ def add_wandb_env_vars(job: Union[xm.Job, xm_cluster.ArrayJob]):
         job = copy.copy(job)  # type: ignore
 
         async def job_gen(work_unit: xm_cluster.ClusterWorkUnit):
-            experiment_title = work_unit.experiment._experiment_title
+            experiment_title = work_unit.experiment._experiment_title  # type: ignore
             xid = work_unit.experiment_id
             wid = work_unit.work_unit_id
 

lxm3/experimental/image_cache.py

@@ -157,7 +157,7 @@ def get_cached_image(image_spec: str, cache_dir: str) -> ImageInfo:
 
     elif transport == "podman-daemon":
         try:
-            import podman
+            import podman  # type: ignore
         except ImportError:
             raise ImportError(
                 "podman package is required to use docker daemon transport"

lxm3/xm_cluster/execution/local.py

@@ -47,7 +47,7 @@ def _is_gpu_requested(cls, executor: executors.Local) -> bool:
 
     @classmethod
     def _create_setup_cmds(
-        cls, executable: executables.Command, executor: executors.GridEngine
+        cls, executable: executables.Command, executor: executors.Local
     ) -> str:
         del executor
         cmds = ["echo >&2 INFO[$(basename $0)]: Running on host $(hostname)"]

lxm3/xm_cluster/experiment.py

@@ -49,8 +49,6 @@ async def _launch(
 class ClusterWorkUnit(xm.WorkUnit):
     """A mock version of WorkUnit with abstract methods implemented."""
 
-    experiment: "ClusterExperiment"
-
     def __init__(
         self,
         experiment: "ClusterExperiment",
@@ -72,26 +70,29 @@ def __init__(
     async def _launch_job_group(
         self,
         job_group: xm.JobGroup,
-        args: Optional[Mapping[str, Any]],
+        args_view: Optional[Mapping[str, Any]],
         identity: str,
     ) -> None:
         """Appends the job group to the launched_jobs list."""
         del identity
 
         async with self._work_unit_id_predictor.submit_id(self._work_unit_id):  # type: ignore
-            await self._submit_job_for_execution(job_group, args)
+            await self._submit_job_for_execution(job_group, args_view)
 
-    async def _launch_job_config(self, job, args, identity):
+    async def _launch_job_config(self, job_config, args_view, identity):
         del identity
-        assert not args
+        assert not args_view
         async with self._work_unit_id_predictor.submit_id(self._work_unit_id):  # type: ignore
-            await self._submit_job_for_execution(job, args)
+            assert isinstance(job_config, array_job_lib.ArrayJob)
+            await self._submit_job_for_execution(job_config, args_view)
 
     async def _submit_job_for_execution(
         self, job: Union[xm.JobGroup, array_job_lib.ArrayJob], args
     ):
         launch_result = await _launch(
-            self.experiment._experiment_title, self.experiment_unit_name, job
+            self.experiment._experiment_title,  # type: ignore
+            self.experiment_unit_name,
+            job,
         )
         self._ingest_handles(launch_result)
 
@@ -195,11 +196,10 @@ async def __aexit__(self, exc_type, exc_value, traceback):
 
     @property
     def work_unit_count(self) -> int:
-        return len(self.work_units)
+        return len(self.work_units())
 
-    @property
     def work_units(self):
-        return self._work_units
+        return {i: wu for i, wu in enumerate(self._work_units)}
 
     @property
     def experiment_id(self) -> int:

pyproject.toml

@@ -129,6 +129,9 @@ exclude = [
   "lxm3/xm",
   "lxm3/_vendor",
 ]
+ignore = [
+  "lxm3/xm_cluster/executable_specs.py"
+]
 
 reportMissingImports = true
 reportMissingTypeStubs = false