Merge pull request #406 from ExaWorks/fix_serialization

Fix serialization

Author: hategan

.mypy

@@ -17,3 +17,5 @@ ignore_missing_imports = True
 [mypy-pystache.*]
 ignore_missing_imports = True
 
+[mypy-typing_compat.*]
+ignore_missing_imports = True

docs/api.rst

@@ -87,7 +87,7 @@ terminal:
 
 
 JobExecutor Base Class
-^^^^^^^^^^^^^^
+^^^^^^^^^^^^^^^^^^^^^^
 
 The ``psij.JobExecutor`` class is abstract, but offers concrete static methods
 for registering, fetching, and listing subclasses of itself.
@@ -216,7 +216,7 @@ multiple
     :noindex:
 
 Other Package Contents
-----------------
+----------------------
 
 .. automodule:: psij.exceptions
     :members:

requirements.txt

@@ -2,3 +2,4 @@ filelock~=3.4
 psutil~=5.9
 pystache>=0.6.0
 typeguard~=2.12
+typing-compat

src/psij/__init__.py

@@ -17,15 +17,14 @@
 from .job_state import JobState
 from .job_status import JobStatus
 from .resource_spec import ResourceSpec, ResourceSpecV1
-from .serialize import Export, Import
 from .version import VERSION
 
 __version__ = VERSION
 
 __all__ = [
     'JobExecutor', 'JobExecutorConfig', 'Job', 'JobStatusCallback', 'JobSpec', 'JobAttributes',
     'JobStatus', 'JobState', 'ResourceSpec', 'ResourceSpecV1', 'Launcher', 'SubmitException',
-    'InvalidJobException', 'UnreachableStateException', 'Export', 'Import'
+    'InvalidJobException', 'UnreachableStateException'
 ]
 
 logger = logging.getLogger(__name__)

src/psij/job_attributes.py

@@ -51,8 +51,44 @@ def get_custom_attribute(self, name: str) -> Optional[object]:
             return None
         return self._custom_attributes[name]
 
+    @property
+    def custom_attributes(self) -> Optional[Dict[str, object]]:
+        """Returns a dictionary with the custom attributes."""
+        return self._custom_attributes
+
+    @custom_attributes.setter
+    def custom_attributes(self, attrs: Optional[Dict[str, object]]) -> None:
+        """
+        Sets all custom attributes from the given dictionary.
+
+        Existing custom attributes defined on this instance of `JobAttributes`
+        are removed.
+
+        Parameters
+        ----------
+        attrs
+            A dictionary with the custom attributes to set.
+        """
+        self._custom_attributes = attrs
+
     def __repr__(self) -> str:
         """Returns a string representation of this object."""
         return 'JobAttributes(duration={}, queue_name={}, project_name={}, reservation_id={}, ' \
                'custom_attributes={})'.format(self.duration, self.queue_name, self.project_name,
                                               self.reservation_id, self._custom_attributes)
+
+    def __eq__(self, o: object) -> bool:
+        """
+        Tests if this JobAttributes object is equal to another object.
+
+        The objects are equal if all their properties are equal.
+        """
+        if not isinstance(o, JobAttributes):
+            return False
+
+        for prop_name in ['duration', 'queue_name', 'project_name', 'reservation_id',
+                          'custom_attributes']:
+            if getattr(self, prop_name) != getattr(o, prop_name):
+                return False
+
+        return True

src/psij/job_spec.py

@@ -1,20 +1,16 @@
 from __future__ import annotations
 
-import sys
+import pathlib
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 from typeguard import check_argument_types
 
-from psij.job_attributes import JobAttributes
-from psij.resource_spec import ResourceSpec
-from psij.utils import path_object_to_full_path as o2p
+import psij.resource_spec
+import psij.job_attributes
 
 
-StrOrPath = Union[str, Path]
-
-
-def _to_path(arg: Optional[StrOrPath]) -> Optional[Path]    :
+def _to_path(arg: Optional[Union[str, Path]]) -> Optional[Path]:
     if isinstance(arg, Path):
         return arg
     elif arg is None:
@@ -28,12 +24,24 @@ class JobSpec(object):
     """A class to hold information about the characteristics of a:class:`~psij.Job`."""
 
     def __init__(self, executable: Optional[str] = None, arguments: Optional[List[str]] = None,
-                 directory: Optional[StrOrPath] = None, name: Optional[str] = None,
+                 # For some odd reason, and only in the constructor, if Path is used directly,
+                 # sphinx fails to find the class. Using Path in the getters and setters does not
+                 # appear to trigger a problem.
+                 directory: Optional[Union[str, pathlib.Path]] = None,
+                 name: Optional[str] = None,
                  inherit_environment: bool = True, environment: Optional[Dict[str, str]] = None,
-                 stdin_path: Optional[StrOrPath] = None, stdout_path: Optional[StrOrPath] = None,
-                 stderr_path: Optional[StrOrPath] = None, resources: Optional[ResourceSpec] = None,
-                 attributes: Optional[JobAttributes] = None, pre_launch: Optional[StrOrPath] = None,
-                 post_launch: Optional[StrOrPath] = None, launcher: Optional[str] = None):
+                 stdin_path: Optional[Union[str, pathlib.Path]] = None,
+                 stdout_path: Optional[Union[str, pathlib.Path]] = None,
+                 stderr_path: Optional[Union[str, pathlib.Path]] = None,
+                 # Importing ResourceSpec directly used to work, but for some unclear reason
+                 # sphinx started complaining about finding duplicate ResourceSpec classes,
+                 # psij.resource_spec.ResourceSpec and psij.ResourceSpec, despite the latter not
+                 # being imported.
+                 resources: Optional[psij.resource_spec.ResourceSpec] = None,
+                 attributes: Optional[psij.job_attributes.JobAttributes] = None,
+                 pre_launch: Optional[Union[str, pathlib.Path]] = None,
+                 post_launch: Optional[Union[str, pathlib.Path]] = None,
+                 launcher: Optional[str] = None):
         """
         Constructs a `JobSpec` object while allowing its properties to be initialized.
 
@@ -123,7 +131,8 @@ def __init__(self, executable: Optional[str] = None, arguments: Optional[List[st
         self._stdout_path = _to_path(stdout_path)
         self._stderr_path = _to_path(stderr_path)
         self.resources = resources
-        self.attributes = attributes if attributes is not None else JobAttributes()
+        self.attributes = attributes if attributes is not None else \
+            psij.job_attributes.JobAttributes()
         self._pre_launch = _to_path(pre_launch)
         self._post_launch = _to_path(post_launch)
         self.launcher = launcher
@@ -140,146 +149,88 @@ def name(self) -> Optional[str]:
         else:
             return self._name
 
+    @name.setter
+    def name(self, value: Optional[str]) -> None:
+        self._name = value
+
     @property
     def directory(self) -> Optional[Path]:
+        """The directory, on the compute side, in which the executable is to be run."""
         return self._directory
 
     @directory.setter
-    def directory(self, directory: Optional[StrOrPath]) -> None:
+    def directory(self, directory: Optional[Union[str, Path]]) -> None:
         self._directory = _to_path(directory)
 
     @property
     def stdin_path(self) -> Optional[Path]:
+        """Path to a file whose contents will be sent to the job's standard input."""
         return self._stdin_path
 
     @stdin_path.setter
-    def stdin_path(self, stdin_path: Optional[StrOrPath]) -> None:
+    def stdin_path(self, stdin_path: Optional[Union[str, Path]]) -> None:
         self._stdin_path = _to_path(stdin_path)
 
     @property
     def stdout_path(self) -> Optional[Path]:
+        """A path to a file in which to place the standard output stream of the job."""
         return self._stdout_path
 
     @stdout_path.setter
-    def stdout_path(self, stdout_path: Optional[StrOrPath]) -> None:
+    def stdout_path(self, stdout_path: Optional[Union[str, Path]]) -> None:
         self._stdout_path = _to_path(stdout_path)
 
     @property
     def stderr_path(self) -> Optional[Path]:
+        """A path to a file in which to place the standard error stream of the job."""
         return self._stderr_path
 
     @stderr_path.setter
-    def stderr_path(self, stderr_path: Optional[StrOrPath]) -> None:
+    def stderr_path(self, stderr_path: Optional[Union[str, Path]]) -> None:
         self._stderr_path = _to_path(stderr_path)
 
     @property
     def pre_launch(self) -> Optional[Path]:
+        """
+        An optional path to a pre-launch script.
+
+        The pre-launch script is sourced before the launcher is invoked. It, therefore, runs on
+        the service node of the job rather than on all of the compute nodes allocated to the job.
+        """
         return self._pre_launch
 
     @pre_launch.setter
-    def pre_launch(self, pre_launch: Optional[StrOrPath]) -> None:
+    def pre_launch(self, pre_launch: Optional[Union[str, Path]]) -> None:
         self._pre_launch = _to_path(pre_launch)
 
     @property
     def post_launch(self) -> Optional[Path]:
+        """
+        An optional path to a post-launch script.
+
+        The post-launch script is sourced after all the ranks of the job executable complete and
+        is sourced on the same node as the pre-launch script.
+        """
         return self._post_launch
 
     @post_launch.setter
-    def post_launch(self, post_launch: Optional[StrOrPath]) -> None:
+    def post_launch(self, post_launch: Optional[Union[str, Path]]) -> None:
         self._post_launch = _to_path(post_launch)
 
-    def _init_job_spec_dict(self) -> Dict[str, Any]:
-        """Returns jobspec structure as dict."""
-        # convention:
-        #  - if expected value is a string then the dict is initialized with an empty string
-        # - if the expected value is an object than the key is initialzied with None
-
-        job_spec: Dict[str, Any]
-        job_spec = {
-            'name': '',
-            'executable': '',
-            'arguments': [],
-            'directory': None,
-            'inherit_environment': True,
-            'environment': {},
-            'stdin_path': None,
-            'stdout_path': None,
-            'stderr_path': None,
-            'resources': None,
-            'attributes': None,
-            'launcher': None
-        }
-
-        return job_spec
+    def __eq__(self, o: object) -> bool:
+        """
+        Tests if this JobSpec is equal to another.
 
-    @property
-    def to_dict(self) -> Dict[str, Any]:
-        """Returns a dictionary representation of this object."""
-        d = self._init_job_spec_dict
-
-        # Map properties to keys
-        d['name'] = self.name
-        d['executable'] = self.executable
-        d['arguments'] = self.arguments
-        d['directory'] = o2p(self.directory)
-        d['inherit_environment'] = self.inherit_environment
-        d['environment'] = self.environment
-        d['stdin_path'] = o2p(self.stdin_path)
-        d['stdout_path'] = o2p(self.stdout_path)
-        d['stderr_path'] = o2p(self.stderr_path)
-        d['resources'] = self.resources
-
-        # Handle attributes property
-        if self.attributes:
-            d['attributes'] = {
-                'duration': '',
-                'queue_name': '',
-                'project_name': '',
-                'reservation_id': '',
-                'custom_attributes': {},
-            }
-            for k, v in self.attributes.__dict__.items():
-                if k in ['duration', 'queue_name', 'project_name', 'reservation_id']:
-                    if v:
-                        d['attributes'][k] = str(v)
-                    else:
-                        d['attributes'][k] = v
-                elif k == "_custom_attributes":
-                    if v:
-                        for ck, cv in v.items():
-                            if not type(cv).__name__ in ['str',
-                                                         'list',
-                                                         'dict',
-                                                         'NoneType',
-                                                         'bool',
-                                                         'int']:
-                                sys.stderr.write("Unsupported type "
-                                                 + type(cv).__name__
-                                                 + " in JobAttributes.custom_attributes for key "
-                                                 + ck
-                                                 + ", skipping\n")
-                        else:
-                            if ck:
-                                d['attributes']['custom_attributes'][ck] = str(cv)
-                            else:
-                                d['attributes']['custom_attributes'][ck] = cv
-                else:
-                    sys.stderr.write("Unsupported attribute " + k + ", skipping attribute\n")
-        else:
-            d['attributes'] = None
-
-        if self.resources:
-
-            d['resources'] = {
-                'node_count': None,
-                'process_count': None,
-                'process_per_node': None,
-                'cpu_cores_per_process': None,
-                'gpu_cores_per_process': None,
-                'exclusive_node_use': None
-            }
-            r = self.resources.__dict__
-            for k in d['resources'].keys():
-                d['resources'][k] = r[k] if k in r else None
-
-        return d
+        Two job specifications are equal if they represent the same job. That is, if all
+        properties are pair-wise equal.
+        """
+        if not isinstance(o, JobSpec):
+            return False
+
+        for prop_name in ['name', 'executable', 'arguments', 'directory', 'inherit_environment',
+                          'environment', 'stdin_path', 'stdout_path', 'stderr_path', 'resources',
+                          'attributes', 'pre_launch', 'post_launch', 'launcher']:
+            if getattr(self, prop_name) != getattr(o, prop_name):
+                return False
+
+        return True

src/psij/resource_spec.py

@@ -28,6 +28,22 @@ def version(self) -> int:
         """Returns the version of this resource specification class."""
         pass
 
+    @staticmethod
+    def get_instance(version: int) -> 'ResourceSpec':
+        """
+        Creates an instance of a `ResourceSpec` of the specified version.
+
+        Parameters
+        ----------
+        version
+            The version of `ResourceSpec` to instantiate. For example, if `version == 1`, this
+            method will return a new instance of `ResourceSpecV1`.
+        """
+        if version == 1:
+            return ResourceSpecV1()
+        else:
+            raise ValueError()
+
 
 class ResourceSpecV1(ResourceSpec):
     """This class implements V1 of the PSI/J resource specification."""
@@ -168,3 +184,20 @@ def computed_processes_per_node(self) -> int:
     def version(self) -> int:
         """Returns the version of this `ResourceSpec`, which is 1 for this class."""
         return 1
+
+    def __eq__(self, o: object) -> bool:
+        """
+        Tests if this ResourceSpecV1 is equal to another object.
+
+        The objects are equal if all their properties are equal.
+        """
+        if not isinstance(o, ResourceSpecV1):
+            return False
+
+        for prop_name in ['node_count', 'process_count', 'processes_per_node',
+                          'cpu_cores_per_process', 'gpu_cores_per_process',
+                          'exclusive_node_use']:
+            if getattr(self, prop_name) != getattr(o, prop_name):
+                return False
+
+        return True