Update typehints

Signed-off-by: Kai Fricke <[email protected]>

python/ray/tune/trainable/trainable.py

@@ -1371,7 +1371,7 @@ def step(self):
  """
  raise NotImplementedError
 
- def save_checkpoint(self, checkpoint_dir: str) -> Optional[Union[str, Dict]]:
+ def save_checkpoint(self, checkpoint_dir: str) -> Optional[Dict]:
  """Subclasses should override this to implement ``save()``.
 
  Warning:
@@ -1395,11 +1395,9 @@ def save_checkpoint(self, checkpoint_dir: str) -> Optional[Union[str, Dict]]:
  the provided path may be temporary and moved.
 
  Returns:
- A dict or string. If string, the return value is expected to be
- the `checkpoint_dir`. If dict, the return value will
- be automatically serialized by Tune. In both cases, the return value
- is exactly what will be passed to ``Trainable.load_checkpoint()``
- upon restore.
+ A dict or None. If dict, the return value will
+ be automatically serialized by Tune. In that case,
+ ``Trainable.load_checkpoint()`` will receive the dict upon restore.
 
  Example:
  >>> trainable, trainable1, trainable2 = ... # doctest: +SKIP
@@ -1412,7 +1410,7 @@ def save_checkpoint(self, checkpoint_dir: str) -> Optional[Union[str, Dict]]:
  """
  raise NotImplementedError
 
- def load_checkpoint(self, checkpoint: Union[Dict, str]):
+ def load_checkpoint(self, checkpoint: Optional[Dict]):
  """Subclasses should override this to implement restore().
 
  Warning:
@@ -1462,10 +1460,8 @@ def load_checkpoint(self, checkpoint: Union[Dict, str]):
 
  Args:
  checkpoint: If dict, the return value is as
- returned by `save_checkpoint`. If a string, then it is
- a checkpoint path that may have a different prefix than that
- returned by `save_checkpoint`. The directory structure
- underneath the `checkpoint_dir` from `save_checkpoint` is preserved.
+ returned by ``save_checkpoint``. Otherwise, the directory
+ the checkpoint was stored in.
  """
  raise NotImplementedError