def __init__(self,
name: str,
factories_num: int,
key: str,
username: str,
setup_cmds: Optional[List[str]] = None) -> None:
super().__init__()
self.name = name
self.factories_num = factories_num
self.setup_cmds = setup_cmds
self.key = os.path.abspath(os.path.expanduser(key))
self.username = username
self.debug = is_dev_mode()
self.orchestrator: Optional[OrchestratorInstance] = None
self.factories: List[FactoryInsT] = []
def main(args: argparse.Namespace) -> None:
"""Execute command based on given config"""
if is_dev_mode():
print(cl.RA(ASCII_LOGO_DEV))
print(cl.BL(f"Location: {get_flambe_repo_location()}\n"))
else:
print(cl.RA(ASCII_LOGO))
print(cl.BL(f"VERSION: {flambe.__version__}\n"))
# Pass original module for ray / pickle
make_component(torch.nn.Module, TORCH_TAG_PREFIX, only_module='torch.nn')
# torch.optim.Optimizer exists, ignore mypy
make_component(
torch.optim.Optimizer,
TORCH_TAG_PREFIX, # type: ignore
only_module='torch.optim')
make_component(torch.optim.lr_scheduler._LRScheduler,
TORCH_TAG_PREFIX,
only_module='torch.optim.lr_scheduler')
make_component(ray.tune.schedulers.TrialScheduler, TUNE_TAG_PREFIX)
make_component(ray.tune.suggest.SearchAlgorithm, TUNE_TAG_PREFIX)
# TODO check first if there is cluster as if there is there
# is no need to install extensions
check_system_reqs()
with SafeExecutionContext(args.config) as ex:
if args.cluster is not None:
with SafeExecutionContext(args.cluster) as ex_cluster:
cluster, _ = ex_cluster.preprocess(
secrets=args.secrets, install_ext=args.install_extensions)
runnable, extensions = ex.preprocess(import_ext=False,
secrets=args.secrets)
cluster.run(force=args.force)
if isinstance(runnable, ClusterRunnable):
cluster = cast(Cluster, cluster)
# This is independant to the type of ClusterRunnable
destiny = os.path.join(cluster.get_orch_home_path(),
"extensions")
# Before sending the extensions, they need to be
# downloaded (locally).
t = os.path.join(FLAMBE_GLOBAL_FOLDER, "extensions")
extensions = download_extensions(extensions, t)
# At this point, all remote extensions
# (except pypi extensions)
# have local paths.
new_extensions = cluster.send_local_content(extensions,
destiny,
all_hosts=True)
new_secrets = cluster.send_secrets()
# Installing the extensions is crutial as flambe
# will execute without '-i' flag and therefore
# will assume that the extensions are installed
# in the orchestrator.
cluster.install_extensions_in_orchestrator(new_extensions)
logger.info(cl.GR("Extensions installed in Orchestrator"))
runnable.setup_inject_env(cluster=cluster,
extensions=new_extensions,
force=args.force)
cluster.execute(runnable, new_extensions, new_secrets,
args.force)
else:
raise ValueError(
"Only ClusterRunnables can be executed in a cluster.")
else:
runnable, _ = ex.preprocess(secrets=args.secrets,
install_ext=args.install_extensions)
runnable.run(force=args.force, verbose=args.verbose)