import codecs
import os
import shutil
import sys
import time
import zipfile
from argparse import ArgumentParser
from datetime import datetime
from tempfile import TemporaryDirectory
from typing import *
from bson import ObjectId
from .config import Config, ConfigLoader, config_to_dict, config_defaults
from .events import EventHost, Event
from .formatting import format_as_asctime
from .mlstorage import MLStorageClient, ExperimentDoc
from .utils import (NOT_SET, make_dir_archive, json_dumps, json_loads,
ContextStack, deprecated_arg)
from .utils.remote_doc import merge_updates_into_doc
from .type_check import DiscardMode
from .typing_ import TConfig
__all__ = ['Experiment', 'get_active_experiment']
class Experiment(Generic[TConfig]):
"""
Class to manage the configuration and results of an experiment.
Basic Usage
===========
To use this class, you first define your experiment config with
:class:`mltk.Config`, and then wrap your main experiment routine with an
:class:`Experiment` context, for example, you may write your `main.py` as::
import numpy as np
from mltk import Config, Experiment
class YourConfig(Config):
max_epoch: int = 100
learning_rate: float = 0.001
...
if __name__ == '__main__':
with Experiment(YourConfig) as exp:
# `exp.config` is the configuration values
print('Max epoch: ', exp.config.max_epoch)
print('Learning rate: ', exp.config.learning_rate)
# `exp.output_dir` is the output directory of this experiment
print('Output directory: ', exp.output_dir)
# save result metrics into `exp.output_dir + "/result.json"`
exp.update_results({'test_acc': ...})
# write result arrays into `exp.output_dir + "/data.npz"`
output_file = exp.abspath('data.npz')
np.savez(output_file, predict=...)
Then you may execute your python file via::
python main.py # use the default config to run the file
python main.py --max_epoch=200 # override some of the config
The output directory (i.e., `exp.output_dir`) will be by default
`"./results/" + script_name`, where `script_name` is the file name of the
main module (excluding ".py"). The configuration values will be saved
into `exp.output_dir + "/config.json"`. If the config file already exists,
it will be loaded and merged with the config values specified by the CLI
arguments. This behavior can allow resuming from an interrupted experiment.
If the Python script is executed via `mlrun`, then the output directory
will be assigned by MLStorage server. For example::
mlrun -s http://<server>:<port> -- python main.py --max_epoch=200
You may also get the server URI and the experiment ID assigned by the
server via the properties `id` and `client`.
To resume from an interrupted experiment with `mlrun`::
mlrun --resume-from=<experiment id> -- python main.py
Since `mlrun` will pick up the output directory of the previous experiment,
:class:`Experiment` will correctly restore the configuration values from
`exp.output_dir + "/config.json"`, thus no need to specify the CLI arguments
once again when resuming from an experiment.
"""
def __init__(self,
config_or_cls: Union[Type[TConfig], TConfig],
*,
script_name: str = NOT_SET,
output_dir: Optional[str] = NOT_SET,
args: Optional[Sequence[str]] = NOT_SET,
auto_load_config: bool = True,
auto_save_config: bool = True,
discard_undefind_config_fields: Union[str, DiscardMode] = DiscardMode.WARN,
):
"""
Construct a new :class:`Experiment`.
Args:
config_or_cls: The configuration object, or class.
script_name: The script name. By default use the file name
of ``sys.modules['__main__']`` (excluding ".py").
output_dir: The output directory. If not specified, use
`"./results/" + script_name`, or assigned by MLStorage
server if the experiment is launched by `mlrun`.
Set to :obj:`None` will disable using an output directory.
In this case, any method that attempt to operate on the
output directory will cause an error.
args: The CLI arguments. If not specified, use ``sys.argv[1:]``.
Specifying :obj:`None` will disable parsing the arguments.
auto_load_config: Whether or not to restore configuration
values from the previously saved `output_dir + "/config.json"`?
If `output_dir` is None, this argument will be ignored.
auto_save_config: Whether or not to save configuration
values into `output_dir + "/config.json"`?
If `output_dir` is None, this argument will be ignored.
discard_undefind_config_fields: The mode to deal with undefined
config fields when loading from previously saved config file.
Defaults to ``DiscardMode.WARN``, where the undefined config
fields are automatically discarded, with a warning generated.
"""
# validate the arguments
config_or_cls_okay = True
config = None
if isinstance(config_or_cls, type):
if not issubclass(config_or_cls, Config):
config_or_cls_okay = False
else:
config = config_or_cls()
else:
if not isinstance(config_or_cls, Config):
config_or_cls_okay = False
else:
config = config_or_cls
if not config_or_cls_okay:
raise TypeError(f'`config_or_cls` is neither a Config class, '
f'nor a Config instance: {config_or_cls!r}')
if script_name is NOT_SET:
script_name = os.path.splitext(
os.path.basename(sys.modules['__main__'].__file__))[0]
if output_dir is NOT_SET:
output_dir = os.environ.get('MLSTORAGE_OUTPUT_DIR', NOT_SET)
if output_dir is NOT_SET:
candidate_dir = f'./results/{script_name}'
while os.path.exists(candidate_dir):
time.sleep(0.01)
suffix = format_as_asctime(
datetime.now(),
datetime_format='%Y-%m-%d_%H-%M-%S',
datetime_msec_sep='_',
)
candidate_dir = f'./results/{script_name}_{suffix}'
output_dir = candidate_dir
if output_dir is not None:
output_dir = os.path.abspath(output_dir)
if args is NOT_SET:
args = sys.argv[1:]
if args is not None:
args = tuple(map(str, args))
# memorize the arguments
self._config = config
self._script_name = script_name
self._output_dir = output_dir
self._args = args
self._auto_load_config = auto_load_config
self._auto_save_config = auto_save_config
self._discard_undefind_config_fields = discard_undefind_config_fields
# the event
self._events = EventHost()
self._on_enter = self.events['on_enter']
self._on_exit = self.events['on_exit']
# initialize the MLStorage client if environment variable is set
id = os.environ.get('MLSTORAGE_EXPERIMENT_ID', None) or None
if id is not None:
id = ObjectId(id)
if os.environ.get('MLSTORAGE_SERVER_URI', None):
client = MLStorageClient(os.environ['MLSTORAGE_SERVER_URI'])
else:
client = None
self._id = id
self._client = client
# construct the experiment document
self._remote_doc = ExperimentDoc(client=client, id=id)
self._remote_doc.on_before_push.do(self._on_before_push_doc)
@property
def id(self) -> Optional[str]:
"""Get the experiment ID, if the environment variable is set."""
return self._id
@property
def client(self) -> Optional[MLStorageClient]:
"""Get the MLStorage client, if the environment variable is set."""
return self._client
@property
def doc(self) -> ExperimentDoc:
"""Get the experiment document object."""
return self._remote_doc
@property
def config(self) -> TConfig:
"""
Get the config object.
If you would like to modify this object, you may need to manually call
:meth:`save_config()`, in order to save the modifications to disk.
"""
return self._config
@property
def results(self) -> Optional[Mapping[str, Any]]:
"""
Get the results, or :obj:`None` if no result has been written.
"""
return self.doc.get('result', None)
@property
def script_name(self) -> str:
"""Get the script name of this experiment."""
return self._script_name
@property
def output_dir(self) -> Optional[str]:
"""
Get the output directory of this experiment.
Returns:
The output directory, or :obj:`None` if the experiment object
does not have an output directory.
"""
return self._output_dir
@property
def args(self) -> Optional[Tuple[str]]:
"""Get the CLI arguments of this experiment."""
return self._args
@property
def events(self) -> EventHost:
"""Get the event host."""
return self._events
@property
def on_enter(self) -> Event:
"""
Get the on enter event.
Callback function type: `() -> None`
This event will be triggered when entering an experiment context::
with Experiment(...) as exp:
# this event will be triggered after entering the context,
# and before the following statements
...
"""
return self._on_enter
@property
def on_exit(self) -> Event:
"""
Get the on exit event.
Callback function type: `() -> None`
This event will be triggered when exiting an experiment context::
with Experiment(...) as exp:
...
# this event will be triggered after the above statements,
# and before exiting the context
"""
return self._on_exit
def _require_output_dir(self):
if self._output_dir is None:
raise RuntimeError('No output directory is configured.')
def _on_before_push_doc(self, updates: Dict[str, Any]):
"""Callback to save the updated fields to local JSON file."""
# gather updated fields of: 'config', 'default_config', 'result'
updated_fields = {}
merge_updates_into_doc(
updated_fields,
updates,
keys_to_expand=('config', 'default_config', 'result'))
# if no output dir, do not save anything
if self._output_dir is None:
return
# now save the interested fields
def save_json_file(path, obj, merge: bool = False):
if merge:
try:
if os.path.exists(path):
with codecs.open(path, 'rb', 'utf-8') as f:
old_obj = json_loads(f.read())
obj = merge_updates_into_doc(old_obj, obj)
except Exception: # pragma: no cover
raise IOError(f'Cannot load the existing JSON file: {path}')
obj_json = json_dumps(obj, no_dollar_field=True)
with codecs.open(path, 'wb', 'utf-8') as f:
f.write(obj_json)
# NOTE: Do not makedirs until we've confirmed that some fields need
# to be saved.
if any(key in updated_fields
for key in ('config', 'default_config', 'result')):
os.makedirs(self.output_dir, exist_ok=True)
if 'config' in updated_fields:
save_json_file(os.path.join(self.output_dir, 'config.json'),
updated_fields['config'])
if 'default_config' in updated_fields:
save_json_file(os.path.join(self.output_dir, 'config.defaults.json'),
updated_fields['default_config'])
if 'result' in updated_fields:
save_json_file(os.path.join(self.output_dir, 'result.json'),
updated_fields['result'], merge=True)
def save_config(self):
"""Save the config into local file and to remote server."""
self._remote_doc.update(
{
'config': config_to_dict(self.config, flatten=True),
'default_config': config_to_dict(
config_defaults(self.config), flatten=True)
},
immediately=True
)
def update_results(self,
results: Optional[Dict[str, Any]] = None,
**kwargs):
"""
Update the result dict.
Args:
results: The dict of updates.
**kwargs: The named arguments of updates.
"""
results = dict(results or ())
results.update(kwargs)
self._remote_doc.update({'result': results})
def abspath(self, relpath: str) -> str:
"""
Get the absolute path of a relative path in `output_dir`.
Args:
relpath: The relative path.
Returns:
The absolute path of `relpath`.
"""
self._require_output_dir()
return os.path.join(self.output_dir, relpath)
def make_dirs(self, relpath: str, exist_ok: bool = True) -> str:
"""
Create a directory (and its ancestors if necessary) in `output_dir`.
Args:
relpath: The relative path of the directory.
exist_ok: If :obj:`True`, will not raise error if the directory
already exists.
Returns:
The absolute path of `relpath`.
"""
self._require_output_dir()
path = self.abspath(relpath)
os.makedirs(path, exist_ok=exist_ok)
return path
def make_parent(self, relpath: str, exist_ok: bool = True) -> str:
"""
Create the parent directory of `relpath` (and its ancestors if
necessary) in `output_dir`.
Args:
relpath: The relative path of the entry, whose parent and
ancestors are to be created.
exist_ok: If :obj:`True`, will not raise error if the parent
directory already exists.
Returns:
The absolute path of `relpath`.
"""
self._require_output_dir()
path = self.abspath(relpath)
parent_dir = os.path.split(path)[0]
os.makedirs(parent_dir, exist_ok=exist_ok)
return path
def open_file(self, relpath: str, mode: str, encoding: Optional[str] = None,
make_parent: bool = NOT_SET):
"""
Open a file at `relpath` in `output_dir`.
Args:
relpath: The relative path of the file.
mode: The open mode.
encoding: The text encoding. If not specified, will open the file
in binary mode; otherwise will open it in text mode.
make_parent: If :obj:`True`, will create the parent (and all
ancestors) of `relpath` if necessary. By default, will
create the parent if open the file by writable mode.
Returns:
The opened file.
"""
self._require_output_dir()
if make_parent is NOT_SET:
make_parent = any(s in mode for s in 'aw+')
if make_parent:
path = self.make_parent(relpath)
else:
path = self.abspath(relpath)
if encoding is None:
return open(path, mode)
else:
return codecs.open(path, mode, encoding)
def put_file_content(self,
relpath: str,
content: Union[bytes, str],
append: bool = False,
encoding: Optional[str] = None):
"""
Save content into a file.
Args:
relpath: The relative path of the file.
content: The file content. Must be bytes if `encoding` is not
specified, while text if `encoding` is specified.
append: Whether or not to append to the file?
encoding: The text encoding.
"""
self._require_output_dir()
with self.open_file(relpath, 'ab' if append else 'wb',
encoding=encoding) as f:
f.write(content)
def get_file_content(self, relpath: str, encoding: Optional[str] = None
) -> Union[bytes, str]:
"""
Get the content of a file.
Args:
relpath: The relative path of a file.
encoding: The text encoding. If specified, will decode the
file content using this encoding.
Returns:
The file content.
"""
self._require_output_dir()
with self.open_file(relpath, 'rb', encoding=encoding) as f:
return f.read()
def make_archive(self,
source_dir: str,
archive_file: Optional[str] = None,
delete_source: bool = True):
"""
Pack a directory into a zip archive.
For repeated experiments, pack some result directories into zip
archives will reduce the total inode count of the file system.
Args:
source_dir: The relative path of the source directory.
archive_file: The relative path of the zip archive.
If not specified, will use `source_dir + ".zip"`.
delete_source: Whether or not to delete `source_dir` after
the zip archive has been created?
Returns:
The absolute path of the archive file.
"""
self._require_output_dir()
def _copy_dir(src: str, dst: str):
os.makedirs(dst, exist_ok=True)
for name in os.listdir(src):
f_src = os.path.join(src, name)
f_dst = os.path.join(dst, name)
if os.path.isdir(f_src):
_copy_dir(f_src, f_dst)
else:
shutil.copyfile(f_src, f_dst, follow_symlinks=False)
source_dir = self.abspath(source_dir)
if not os.path.isdir(source_dir):
raise IOError(f'Not a directory: {source_dir}')
if archive_file is None:
archive_file = source_dir.rstrip('/\\') + '.zip'
else:
archive_file = self.abspath(archive_file)
def prepare_parent():
parent_dir = os.path.dirname(archive_file)
os.makedirs(parent_dir, exist_ok=True)
# if the archive already exists, extract it, merge the contents
# in `source_dir` with the extracted files, and then make archive.
if os.path.isfile(archive_file):
with TemporaryDirectory() as temp_dir:
# extract the original zip
with zipfile.ZipFile(archive_file, 'r') as zf:
zf.extractall(temp_dir)
# merge the content
_copy_dir(source_dir, temp_dir)
# make the destination archive
prepare_parent()
make_dir_archive(temp_dir, archive_file)
# otherwise pack the zip archive directly
else:
prepare_parent()
make_dir_archive(source_dir, archive_file)
# now delete the source directory
if delete_source:
shutil.rmtree(source_dir)
return archive_file
def make_archive_on_exit(self,
source_dir: str,
archive_file: Optional[str] = None,
delete_source: bool = True):
"""
Pack a directory into a zip archive when exiting the experiment context.
Args:
source_dir: The relative path of the source directory.
archive_file: The relative path of the zip archive.
If not specified, will use `source_dir + ".zip"`.
delete_source: Whether or not to delete `source_dir` after
the zip archive has been created?
See Also:
:meth:`make_archive()`
"""
self._require_output_dir()
self.on_exit.do(lambda: self.make_archive(
source_dir=source_dir,
archive_file=archive_file,
delete_source=delete_source
))
def __enter__(self) -> 'Experiment[TConfig]':
config_loader = ConfigLoader(self.config)
# build the argument parser
if self.args is not None:
arg_parser = ArgumentParser()
arg_parser.add_argument(
'--output-dir', help='Set the experiment output directory.',
default=NOT_SET, metavar='PATH'
)
arg_parser = config_loader.build_arg_parser(arg_parser)
parsed_args = arg_parser.parse_args(self.args)
output_dir = parsed_args.output_dir
if output_dir is not NOT_SET:
# special hack: override `output_dir` if specified
self._output_dir = os.path.abspath(output_dir)
parsed_args.output_dir = NOT_SET
parsed_args = {
key: value for key, value in vars(parsed_args).items()
if value is not NOT_SET
}
else:
parsed_args = {}
# load previously saved configuration
if self._output_dir is not None and self._auto_load_config:
config_files = [
os.path.join(self.output_dir, 'config.yml'),
os.path.join(self.output_dir, 'config.json'),
]
for config_file in config_files:
try:
if os.path.exists(config_file):
config_loader.load_file(config_file)
except Exception: # pragma: no cover
raise IOError(f'Failed to load config file: '
f'{config_file!r}')
# load the cli arguments
config_loader.load_object(parsed_args)
# finally, generate the config object
self._config = config_loader.get(
discard_undefined=self._discard_undefind_config_fields)
# prepare for the output dir
if self._output_dir is not None:
os.makedirs(self.output_dir, exist_ok=True)
# start the remote doc background worker
self._remote_doc.start_worker()
# save the configuration
if self._auto_save_config:
self.save_config()
# trigger the on enter event
self.on_enter.fire()
# add to context stack
_experiment_stack.push(self)
return self
def __exit__(self, exc_type, exc_val, exc_tb):
try:
self.on_exit.fire()
finally:
if _experiment_stack.top() == self:
_experiment_stack.pop()
self._remote_doc.stop_worker()
def get_active_experiment() -> Optional[Experiment]:
"""
Get the current active experiment object at the top of context stack.
Returns:
The active experiment object.
"""
return _experiment_stack.top()
_experiment_stack: ContextStack[Experiment] = ContextStack()