Merge pull request #315 from ExaWorks/fix/issue_314

add environment type check to `JobSpec`

Author: hategan

requirements.txt

@@ -1,3 +1,4 @@
 filelock
 psutil
 pystache>=0.6.0
+typeguard

src/psij/executors/batch/batch_scheduler_executor.py

@@ -194,9 +194,9 @@ def submit(self, job: Job) -> None:
         """See :func:`~psij.JobExecutor.submit`."""
         logger.info('Job %s: submitting', job.id)
         self._ensure_work_dir()
-        assert (job.spec)
 
-        job.executor = self
+        self._check_job(job)
+
         context = self._create_script_context(job)
 
         # assumes job ids are unique
@@ -551,7 +551,7 @@ def list(self) -> List[str]:
         Implementations are encouraged to restrict the results to jobs accessible by the current
         user.
         """
-        pass
+        raise NotImplementedError()
 
 
 class _QueuePollThread(Thread):

src/psij/executors/flux.py

@@ -124,14 +124,14 @@ def _add_flux_callbacks(self, job: Job, fut: flux.job.FluxExecutorFuture) -> Non
 
     def submit(self, job: Job) -> None:
         """See :func:`~psij.job_executor.JobExecutor.submit`."""
-        assert job.spec
-        assert job.spec.attributes
-        job.executor = self
-        if isinstance(job.spec.resources, ResourceSpecV1):
-            resources = job.spec.resources
-        elif isinstance(job.spec.resources, ResourceSpec):
+        spec = self._check_job(job)
+
+        assert spec.attributes
+        if isinstance(spec.resources, ResourceSpecV1):
+            resources = spec.resources
+        elif isinstance(spec.resources, ResourceSpec):
             raise InvalidJobException(
-                f"ResourceSpec version {job.spec.resources.version} not supported"
+                f"ResourceSpec version {spec.resources.version} not supported"
             )
         else:
             resources = ResourceSpecV1(process_count=1, cpu_cores_per_process=1)
@@ -142,24 +142,24 @@ def submit(self, job: Job) -> None:
             )
         if resources.processes_per_node:
             raise InvalidJobException("Flux does not support processes_per_node")
-        if not job.spec.executable:
+        if not spec.executable:
             raise InvalidJobException("Job must have an executable")
-        argv = list(job.spec.arguments) if job.spec.arguments else []
-        argv.insert(0, job.spec.executable)
+        argv = list(spec.arguments) if spec.arguments else []
+        argv.insert(0, spec.executable)
         flux_jobspec = flux.job.JobspecV1.from_command(
             argv,
             num_tasks=resources.process_count,
             cores_per_task=resources.cpu_cores_per_process,
             gpus_per_task=resources.gpu_cores_per_process,
             num_nodes=resources.node_count,
         )
-        if job.spec.stdout_path:
-            flux_jobspec.stdout = job.spec.stdout_path
-        if job.spec.stdin_path:
-            flux_jobspec.stdin = job.spec.stdin_path
-        if job.spec.stderr_path:
-            flux.jobspec.stderr = job.spec.stderr_path
-        flux_jobspec.duration = job.spec.attributes.duration.total_seconds()
+        if spec.stdout_path:
+            flux_jobspec.stdout = spec.stdout_path
+        if spec.stdin_path:
+            flux_jobspec.stdin = spec.stdin_path
+        if spec.stderr_path:
+            flux.jobspec.stderr = spec.stderr_path
+        flux_jobspec.duration = spec.attributes.duration.total_seconds()
         fut = self._flux_executor.submit(flux_jobspec)
         self._add_flux_callbacks(job, fut)
 

src/psij/executors/local.py

@@ -193,10 +193,7 @@ def submit(self, job: Job) -> None:
 
         :param job: The job to be submitted.
         """
-        spec = job.spec
-        if not spec:
-            raise InvalidJobException('Missing specification')
-        job.executor = self
+        spec = self._check_job(job)
 
         p = _ChildProcessEntry(job, self, self._get_launcher(self._get_launcher_name(spec)))
         assert p.launcher

src/psij/executors/rp.py

@@ -106,11 +106,8 @@ def submit(self, job: Job) -> None:
 
         :param job: The job to be submitted.
         """
-        spec = job.spec
-        if not spec:
-            raise InvalidJobException('Missing specification')
+        self._check_job(job)
 
-        job.executor = self
         try:
             td = self._job_2_descr(job)
             task = self._tmgr.submit_tasks(td)

src/psij/job_attributes.py

@@ -1,6 +1,8 @@
 from datetime import timedelta
 from typing import Optional, Dict
 
+from typeguard import check_argument_types
+
 
 class JobAttributes(object):
     """A class containing ancillary job information that describes how a job is to be run."""
@@ -27,6 +29,8 @@ def __init__(self, duration: timedelta = timedelta(minutes=10),
             :class:`~psij.JobExecutor` define and are responsible for interpreting custom
             attributes.
         """
+        assert check_argument_types()
+
         self.duration = duration
         self.queue_name = queue_name
         self.project_name = project_name

src/psij/job_executor.py

@@ -5,12 +5,14 @@
 from typing import Optional, Dict, List, Type, cast, Union, Callable, Set
 
 import psij
+from psij import InvalidJobException
 from psij.descriptor import Descriptor, _VersionEntry
 from psij._plugins import _register_plugin, _get_plugin_class, _print_plugin_status
 from psij.job_status import JobStatus
 from psij.job import Job, JobStatusCallback, FunctionJobStatusCallback
 from psij.job_executor_config import JobExecutorConfig
 from psij.job_launcher import Launcher
+from psij.job_spec import JobSpec
 
 
 logger = logging.getLogger(__name__)
@@ -57,6 +59,52 @@ def version(self) -> Version:
         """Returns the version of this executor."""
         return cast(Version, getattr(self.__class__, '_VERSION_'))
 
+    def _check_job(self, job: Job) -> JobSpec:
+        """
+        Checks a job for consistency and correctness.
+
+        Verifies that various aspects of the job are correctly specified. This includes precisely
+        the following checks:
+        * the job has a non-null specification
+        * job.spec.environment is a Dict[str, str]
+
+        While this method makes a fair attempt at ensuring the validity of the job, it makes no
+        such guarantees. Specifically, if an executor implementation requires checks not listed
+        above, it should implement them explicitly.
+
+        These checks are meant to trigger common runtime type errors somewhat early and with clear
+        error messages. In production software, these checks can be disabled by invoking Python
+        with one of the optimization flags (`-O` or `-OO`).
+
+        Upon completion, this method sets the :attr:`~psij.Job.executor` attribute of the job and
+        returns the job specification.
+
+        Parameters
+        ----------
+        job
+            The job to validate
+
+        Returns
+        -------
+            A non-null job specification
+        """
+        spec = job.spec
+        if not spec:
+            raise InvalidJobException('Missing specification')
+
+        if __debug__:
+            if spec.environment is not None:
+                for k, v in spec.environment.items():
+                    if not isinstance(k, str):
+                        raise TypeError('environment key "%s" is not a string (%s)'
+                                        % (k, type(k).__name__))
+                    if not isinstance(v, str):
+                        raise TypeError('environment key "%s" has non-string value (%s)'
+                                        % (k, type(v).__name__))
+
+        job.executor = self
+        return spec
+
     @abstractmethod
     def submit(self, job: Job) -> None:
         """

src/psij/job_spec.py

@@ -1,6 +1,9 @@
 import sys
 from pathlib import Path
 from typing import Optional, List, Dict, Any
+
+from typeguard import check_argument_types
+
 from psij.job_attributes import JobAttributes
 from psij.resource_spec import ResourceSpec
 
@@ -53,6 +56,8 @@ def __init__(self, name: Optional[str] = None, executable: Optional[str] = None,
         :param launcher: The name of a launcher to use, such as "mpirun", "srun", "single", etc.
             For a list of available launchers,:ref:`launchers`
         """
+        assert check_argument_types()
+
         self._name = name
         self.executable = executable
         self.arguments = arguments

src/psij/resource_spec.py

@@ -1,6 +1,8 @@
 from abc import ABC, abstractmethod
 from typing import Optional, List
 
+from typeguard import check_argument_types
+
 from psij.exceptions import InvalidJobException
 
 
@@ -58,6 +60,8 @@ def __init__(self, node_count: Optional[int] = None,
         :param gpu_cores_per_process:
         :param exclusive_node_use:
         """
+        assert check_argument_types()
+
         self.node_count = node_count
         self.process_count = process_count
         self.processes_per_node = processes_per_node

tests/test_job_spec.py

@@ -0,0 +1,39 @@
+
+import pytest
+
+from psij import JobSpec, JobExecutor, Job
+
+
+def _test_spec(spec: JobSpec) -> None:
+    ex = JobExecutor.get_instance('local')
+    ex.submit(Job(spec))
+
+
+def test_environment_types() -> None:
+    with pytest.raises(TypeError):
+        _test_spec(JobSpec(environment={'foo': 1}))  # type: ignore
+
+    with pytest.raises(TypeError):
+        _test_spec(JobSpec(environment={1: 'foo'}))  # type: ignore
+
+    with pytest.raises(TypeError):
+        spec = JobSpec()
+        spec.environment = {'foo': 'bar'}
+        spec.environment['buz'] = 2  # type: ignore
+        _test_spec(spec)
+
+    spec = JobSpec()
+    assert spec.environment is None
+
+    spec = JobSpec(environment={'foo': 'bar'})
+    assert spec.environment['foo'] == 'bar'  # type: ignore
+
+    spec = JobSpec()
+    spec.environment = {'foo': 'bar'}
+    assert spec.environment['foo'] == 'bar'
+
+    spec.environment = {'foo': 'biz'}
+    assert spec.environment['foo'] == 'biz'
+
+
+test_environment_types()