diff --git a/lxm3/contrib/wandb.py b/lxm3/contrib/wandb.py index ac02d7e..3c2768b 100644 --- a/lxm3/contrib/wandb.py +++ b/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 diff --git a/lxm3/experimental/image_cache.py b/lxm3/experimental/image_cache.py index 35ddb54..d2aade4 100644 --- a/lxm3/experimental/image_cache.py +++ b/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" diff --git a/lxm3/xm_cluster/execution/local.py b/lxm3/xm_cluster/execution/local.py index ea57bc6..95398a5 100644 --- a/lxm3/xm_cluster/execution/local.py +++ b/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)"] diff --git a/lxm3/xm_cluster/experiment.py b/lxm3/xm_cluster/experiment.py index c8bb013..eb735f2 100644 --- a/lxm3/xm_cluster/experiment.py +++ b/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: diff --git a/pyproject.toml b/pyproject.toml index a2bf7a1..3c4ae2d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -129,6 +129,9 @@ exclude = [ "lxm3/xm", "lxm3/_vendor", ] +ignore = [ + "lxm3/xm_cluster/executable_specs.py" +] reportMissingImports = true reportMissingTypeStubs = false