def _dummy_init(
*,
distributed: Optional[core.DistributedContext] = None,
# TODO(DET-6153): allow a Union[StorageManager, str] here.
storage_manager: Optional[storage.StorageManager] = None,
preempt_mode: core.PreemptMode = core.PreemptMode.WorkersAskChief,
) -> Context:
"""
Build a core.Context suitable for running off-cluster. This is normally called by init()
when it is detected that there is no ClusterInfo available, but can be invoked directly for
e.g. local test mode.
"""
distributed = distributed or core.DummyDistributedContext()
preempt = core.DummyPreemptContext(distributed, preempt_mode)
if storage_manager is None:
base_path = appdirs.user_data_dir("determined")
logger.info("no storage_manager provided; storing checkpoints in {base_path}")
storage_manager = storage.SharedFSStorageManager(base_path)
checkpoint = core.DummyCheckpointContext(distributed, storage_manager)
train = core.DummyTrainContext()
searcher = core.DummySearcherContext(distributed)
_install_stacktrace_on_sigusr1()
return Context(
distributed=distributed,
checkpoint=checkpoint,
preempt=preempt,
train=train,
searcher=searcher,
)
def manager(tmp_path: pathlib.Path) -> storage.StorageManager:
return storage.SharedFSStorageManager(str(tmp_path))
def init(
*,
distributed: Optional[core.DistributedContext] = None,
# TODO: figure out a better way to deal with checkpointing in the local training case.
storage_manager: Optional[storage.StorageManager] = None,
preempt_mode: core.PreemptMode = core.PreemptMode.WorkersAskChief,
tensorboard_mode: core.TensorboardMode = core.TensorboardMode.AUTO,
) -> Context:
"""
``core.init()`` builds a :class:`core.Context <determined.core.Context>` for use with the Core
API.
Always use ``with core.init() as context`` instead of instantiating a ``core.Context`` directly.
Certain components of the Core API may be configured by passing arguments to ``core.init()``.
The only arg that is required is a ``DistributedContext``, and even that is only required for
for multi-slot tasks.
All of your training must occur within the scope of the ``with core.init() as core_context``, as
there are resources necessary for training which start in the ``core.Context``'s ``__enter__``
method and must be cleaned up in its ``__exit__()`` method.
Arguments:
distributed (``core.DistributedContext``, optional): Passing a ``DistributedContext`` is
required for multi-slot training, but unnecessary for single-slot training. Defaults to
``None``.
preempt_mode (``core.PreemptMode``, optional): Configure the calling pattern for the
``core_context.preempt.should_preempt()`` method. See
:class:`~determined.core.PreemptMode` for more detail. Defaults to ``WorkersAskChief``.
storage_manager: Internal use only.
tensorboard_mode (``core.TensorboardMode``, optional): Define how Tensorboard
metrics and profiling data are retained. See
:class:`~determined.core.TensorboardMode`` for more detail. Defaults to ``AUTO``.
"""
info = det.get_cluster_info()
if info is None:
return _dummy_init(distributed=distributed, storage_manager=storage_manager)
# We are on the cluster.
cert = certs.default_load(info.master_url)
session = Session(info.master_url, None, None, cert, max_retries=get_max_retries_config())
if distributed is None:
if len(info.container_addrs) > 1 or len(info.slot_ids) > 1:
raise ValueError("you must provide a valid DistributedContext for a multi-slot task")
distributed = distributed or core.DummyDistributedContext()
preempt = core.PreemptContext(session, info.allocation_id, distributed, preempt_mode)
# At present, we only support tensorboards in Trial tasks.
tbd_writer = None
train = None
searcher = None
if info.task_type == "TRIAL":
# Prepare the tensorboard hooks.
tensorboard_manager = tensorboard.build(
info.cluster_id,
str(info.trial.experiment_id),
str(info.trial.trial_id),
info.trial._config["checkpoint_storage"],
container_path=constants.SHARED_FS_CONTAINER_PATH,
)
if tensorboard_mode == core.TensorboardMode.AUTO:
tbd_writer = tensorboard.get_metric_writer()
train = core.TrainContext(
session,
info.trial.trial_id,
info.trial._trial_run_id,
info.trial.experiment_id,
distributed,
tensorboard_mode,
tensorboard_manager,
tbd_writer,
)
units = core._parse_searcher_units(info.trial._config)
searcher = core.SearcherContext(
session,
distributed,
info.trial.trial_id,
info.trial._trial_run_id,
info.allocation_id,
units,
)
if storage_manager is None:
storage_manager = storage.build(
info.trial._config["checkpoint_storage"],
container_path=constants.SHARED_FS_CONTAINER_PATH,
)
checkpoint = core.CheckpointContext(
distributed,
storage_manager,
session,
info.task_id,
info.allocation_id,
tensorboard_mode,
tensorboard_manager,
)
else:
# TODO: support checkpointing for non-trial tasks.
if storage_manager is None:
base_path = appdirs.user_data_dir("determined")
logger.info("no storage_manager provided; storing checkpoints in {base_path}")
storage_manager = storage.SharedFSStorageManager(base_path)
checkpoint = core.DummyCheckpointContext(distributed, storage_manager)
_install_stacktrace_on_sigusr1()
return Context(
distributed=distributed,
checkpoint=checkpoint,
preempt=preempt,
train=train,
searcher=searcher,
)