def setUp(self):
self.codalab_manager = CodaLabManager()
self.codalab_manager.config['server']['class'] = 'SQLiteModel'
self.bundle_manager = BundleManager(self.codalab_manager)
self.download_manager = self.codalab_manager.download_manager()
self.upload_manager = self.codalab_manager.upload_manager()
# Create a standard user
self.user_id = generate_uuid()
self.bundle_manager._model.add_user(
"codalab_standard",
"*****@*****.**",
"Test",
"User",
"password",
"Stanford",
user_id=self.user_id,
)
# Create a root user
self.root_user_id = self.codalab_manager.root_user_id()
self.bundle_manager._model.add_user(
"codalab_root",
"*****@*****.**",
"Test",
"User",
"password",
"Stanford",
user_id=self.root_user_id,
)
def __init__(self, args):
self.args = args
self.codalab_manager = CodaLabManager()
self.codalab_client = self.codalab_manager.client(args.server)
self.staged_uuids = []
self.last_worker_start_time = 0
logger.info('Started worker manager.')
def _create_cli(self, worksheet_uuid):
manager = CodaLabManager(
temporary=True,
clients={settings.BUNDLE_SERVICE_URL: self.client})
manager.set_current_worksheet_uuid(self.client, worksheet_uuid)
cli = bundle_cli.BundleCLI(manager, headless=True)
return cli
def main(args):
manager = CodaLabManager()
model = manager.model()
# Get the the message
subject = args.subject
with open(args.body_file) as f:
body_template = f.read()
mime_type = 'html' if args.body_file.endswith('.html') else 'plain'
# Figure out who we want to send
to_send_list = get_to_send_list(model, args.threshold)
sent_list = get_sent_list(args.sent_file)
sent_emails = set(info['email'] for info in sent_list)
pending_to_send_list = [info for info in to_send_list if info['email'] not in sent_emails]
print 'Already sent %d emails, %d to go' % (len(sent_list), len(pending_to_send_list))
for i, info in enumerate(pending_to_send_list):
if args.only_email and args.only_email != info['email']:
continue
# Derived fields
info['greeting_name'] = info['first_name'] or info['last_name'] or info['user_name']
info['full_name'] = ' '.join([x for x in [info['first_name'], info['last_name']] if x])
info['email_description'] = '%s <%s>' % (info['full_name'], info['email']) if info['full_name'] else info['email']
info['sent_time'] = time.time()
print 'Sending %s/%s (%s>=%s, doit=%s): [%s] %s' % \
(i, len(pending_to_send_list), info['notifications'], args.threshold, args.doit, info['user_name'], info['email_description'])
# Apply template to get body of message
body = body_template
for field, value in info.items():
body = body.replace('{{' + field + '}}', unicode(value or ''))
if args.verbose >= 1:
print 'To : %s' % info['email_description']
print 'Subject : %s' % subject
print body
print '-------'
if not args.doit:
continue
# Send the actual email
manager.emailer.send_email(
recipient=info['email_description'],
subject=subject,
body=body,
mime_type=mime_type,
)
# Record that we sent
with open(args.sent_file, 'a') as f:
print >>f, json.dumps(info)
f.flush()
def _create_cli(self, worksheet_uuid):
output_buffer = StringIO()
manager = CodaLabManager(temporary=True,
clients={'local': self.client})
manager.set_current_worksheet_uuid(self.client, worksheet_uuid)
cli = bundle_cli.BundleCLI(manager,
headless=True,
stdout=output_buffer,
stderr=output_buffer)
return cli, output_buffer
def setUp(self):
self.codalab_manager = CodaLabManager()
self.codalab_manager.config['server']['class'] = 'SQLiteModel'
self.bundle_manager = BundleManager(self.codalab_manager)
self.user_id = generate_uuid()
self.bundle_manager._model.add_user(
"codalab",
"*****@*****.**",
"Test",
"User",
"password",
"Stanford",
user_id=self.user_id,
)
def __init__(
self,
docker_image,
initial_command="",
manager=None,
dependencies=[],
bundle_locations={},
verbose=False,
stdout=sys.stdout,
stderr=sys.stderr,
):
# Instantiate a CodaLabManager if one is not passed in
self._manager = manager if manager else CodaLabManager()
self._docker_image = docker_image
self._initial_command = initial_command
InteractiveSession._validate_bundle_locations(bundle_locations,
dependencies)
self._dependencies = dependencies
self._bundle_locations = bundle_locations
self._docker_client = docker.from_env(
timeout=InteractiveSession._MAX_SESSION_TIMEOUT)
self._session_uuid = generate_uuid()
self._verbose = verbose
self._stdout = stdout
self._stderr = stderr
def main():
cli = BundleCLI(CodaLabManager())
try:
cli.do_command(sys.argv[1:])
except KeyboardInterrupt:
print('Terminated by Ctrl-C')
sys.exit(130)
def run_migrations_offline():
"""Run migrations in 'offline' mode.
This configures the context with just a URL
and not an Engine, though an Engine is acceptable
here as well. By skipping the Engine creation
we don't even need a DBAPI to be available.
Calls to context.execute() here emit the given string to the
script output.
"""
manager = CodaLabManager()
url = manager.model().engine.url
context.configure(url=url, target_metadata=target_metadata)
with context.begin_transaction():
context.run_migrations()
def run_migrations_online():
"""Run migrations in 'online' mode.
In this scenario we need to create an Engine
and associate a connection with the context.
"""
manager = CodaLabManager()
engine = manager.model().engine
connection = engine.connect()
context.configure(connection=connection, target_metadata=target_metadata)
try:
with context.begin_transaction():
context.run_migrations()
finally:
connection.close()
def run_migrations_offline():
"""Run migrations in 'offline' mode.
This configures the context with just a URL
and not an Engine, though an Engine is acceptable
here as well. By skipping the Engine creation
we don't even need a DBAPI to be available.
Calls to context.execute() here emit the given string to the
script output.
"""
manager = CodaLabManager()
url = manager.model().engine.url
context.configure(url=url, target_metadata=target_metadata)
with context.begin_transaction():
context.run_migrations()
def test_temp_codalab_manager(self):
manager: CodaLabManager = CodaLabManager(temporary=True)
self.assertEqual(manager.state, {'auth': {}, 'sessions': {}})
manager.save_state()
self.assertFalse(
os.path.exists(manager.state_path),
msg=
'Assert that the current state is not written out to state_path for a temporary CodaLabManager',
)
def find_default_editor():
manager = CodaLabManager()
editor = os.getenv('EDITOR')
if editor:
return editor
# If not yet set, use a sane default.
if sys.platform == 'win32':
editor = 'notepad'
else:
editor = 'vi'
return editor
def run_migrations_online():
"""Run migrations in 'online' mode.
In this scenario we need to create an Engine
and associate a connection with the context.
"""
manager = CodaLabManager()
engine = manager.model().engine
connection = engine.connect()
context.configure(
connection=connection,
target_metadata=target_metadata
)
try:
with context.begin_transaction():
context.run_migrations()
finally:
connection.close()
def _create_cli(self, worksheet_uuid):
"""
Create an instance of the CLI.
The CLI uses JsonApiClient to communicate back to the REST API.
This is admittedly not ideal since now the REST API is essentially
making HTTP requests back to itself. Future potential solutions might
include creating a subclass of JsonApiClient that can reroute HTTP
requests directly to the appropriate Bottle view functions.
"""
output_buffer = StringIO()
rest_client = JsonApiClient(self._rest_url(), lambda: get_user_token())
manager = CodaLabManager(
temporary=True,
config=local.config,
clients={
self._rest_url(): rest_client
})
manager.set_current_worksheet_uuid(self._rest_url(), worksheet_uuid)
cli = bundle_cli.BundleCLI(manager, headless=True, stdout=output_buffer, stderr=output_buffer)
return cli, output_buffer
def test_temp_codalab_manager(self):
"""
A codalab manager with temporary state should initialize its state from an existing
state.json file if it is present.
"""
manager: CodaLabManager = CodaLabManager(temporary=True)
self.assertEqual(manager.state, {'auth': {}, 'sessions': {}})
manager.save_state()
self.assertFalse(
os.path.exists(manager.state_path),
msg=
'Assert that the current state is not written out to state_path for a temporary CodaLabManager',
)
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
'--sleep-time',
help='Number of seconds to wait between successive actions.',
type=int,
default=0.5,
)
args = parser.parse_args()
manager = BundleManager(CodaLabManager())
# Register a signal handler to ensure safe shutdown.
for sig in [signal.SIGTERM, signal.SIGINT, signal.SIGHUP]:
signal.signal(sig, lambda signup, frame: manager.signal())
manager.run(args.sleep_time)
def test_temp_codalab_manager_initialize_state(self):
initial_state: Dict = {
'auth': {
"https://worksheets.codalab.org": {
"token_info": {
"access_token": "secret"
}
}
},
'sessions': {},
}
cache_file = tempfile.NamedTemporaryFile(delete=False)
with open(cache_file.name, "w") as f:
json.dump(initial_state, f)
os.environ["CODALAB_STATE"] = cache_file.name
manager: CodaLabManager = CodaLabManager(temporary=True)
self.assertEqual(manager.state, initial_state)
os.remove(cache_file.name)
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
'--sleep-time',
help='Number of seconds to wait between successive actions.',
type=int,
default=0.5,
)
parser.add_argument(
'--worker-timeout-seconds',
help=
'Number of seconds to wait after a worker check-in before determining a worker is offline',
type=int,
default=60,
)
args = parser.parse_args()
manager = BundleManager(CodaLabManager(), args.worker_timeout_seconds)
# Register a signal handler to ensure safe shutdown.
for sig in [signal.SIGTERM, signal.SIGINT, signal.SIGHUP]:
signal.signal(sig, lambda signup, frame: manager.signal())
manager.run(args.sleep_time)
def create_rest_app(manager=CodaLabManager()):
"""Creates and returns a rest app."""
install(SaveEnvironmentPlugin(manager))
install(CheckJsonPlugin())
install(oauth2_provider.check_oauth())
install(CookieAuthenticationPlugin())
install(UserVerifiedPlugin())
install(PublicUserPlugin())
install(ErrorAdapter())
# Replace default JSON plugin with one that handles datetime objects
# Note: ErrorAdapter must come before JSONPlugin to catch serialization errors
uninstall(JSONPlugin())
install(JSONPlugin(json_dumps=DatetimeEncoder().encode))
# JsonApiPlugin must come after JSONPlugin, to inspect and modify response
# dicts before they are serialized into JSON
install(JsonApiPlugin())
for code in range(100, 600):
default_app().error(code)(error_handler)
root_app = Bottle()
root_app.mount('/rest', default_app())
# Look for templates in codalab-worksheets/views
bottle.TEMPLATE_PATH = [
os.path.join(
os.path.dirname(
os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
'views')
]
# Increase the request body size limit to 8 MiB
bottle.BaseRequest.MEMFILE_MAX = 8 * 1024 * 1024
return root_app
observer = Observer()
observer.schedule(event_handler, path, recursive=True)
observer.start()
try:
while True:
time.sleep(10)
except KeyboardInterrupt:
observer.stop()
observer.join()
else:
from codalab.server.bundle_rpc_server import BundleRPCServer
rpc_server = BundleRPCServer(manager)
rpc_server.serve_forever()
def run_cli():
from codalab.lib.bundle_cli import BundleCLI
cli = BundleCLI(manager)
cli.do_command(sys.argv[1:])
if __name__ == '__main__':
manager = CodaLabManager()
# Either start the server or the client.
try:
if len(sys.argv) > 1 and sys.argv[1] == 'server':
run_server()
else:
run_cli()
except KeyboardInterrupt:
print 'Terminated by Ctrl-C'
sys.exit(130)
class WorkerManager(object):
"""
The abstract class for a worker manager. Different backends like AWS,
Azure, Slurm should override `get_worker_jobs` and `start_worker_job`.
The basic architecture of the WorkerManager is extremely simple: to a
first-order approximation, it simply launches `cl-worker`s as AWS/Azure
batch jobs as long as there are staged bundles.
The simplicity means that we don't need to manage state about how workers
and bundles are related - that logic is complex and is done by the usual
worker system, and we don't want to duplicate that for every possible
worker backend.
More specifically, a worker manager will monitor a job queue to see how
many worker jobs are running, and try to keep that between `min_workers`
and `max_workers`. It will also monitor the staged bundles that satisfy a
certain `search` criterion. If there are staged bundles then it will issue
a `start_worker_job()` call, provided some other conditions are met (e.g.,
don't start workers too fast).
The WorkerManager is all client-side code, so it can be customized as one
sees fit.
Notes:
- The worker manager is not visible via CodaLab (i.e., CodaLab has no
notion of a worker manager or what it's trying to do - all it sees is
bundles and workers). One needs to monitor the AWS/Azure Batch system
separately.
- Resource handling is not currently supported. Generally, the safe thing
is to create a separate queue for different resource needs and put the
burden of deciding on the user.
"""
# Subcommand name to use for this worker manager type
NAME = 'worker-manager'
DESCRIPTION = 'Base class for Worker Managers, please implement for your deployment'
@staticmethod
def add_arguments_to_subparser(subparser):
"""
Add any arguments specific to this worker manager to the given subparser
"""
raise NotImplementedError
def __init__(self, args):
self.args = args
self.codalab_manager = CodaLabManager()
self.codalab_client = self.codalab_manager.client(args.server)
self.staged_uuids = []
self.last_worker_start_time = 0
logger.info('Started worker manager.')
def get_worker_jobs(self):
"""Return a list of `WorkerJob`s."""
raise NotImplementedError
def start_worker_job(self):
"""Start a new `WorkerJob`."""
raise NotImplementedError
def build_command(self, worker_id, work_dir):
command = [
self.args.worker_executable,
'--server',
self.args.server,
'--verbose',
'--exit-when-idle',
'--idle-seconds',
str(self.args.worker_idle_seconds),
'--work-dir',
work_dir,
'--id',
f'$(hostname -s)-{worker_id}',
'--network-prefix',
'cl_worker_{}_network'.format(worker_id),
]
# Additional optional arguments
if self.args.worker_tag:
command.extend(['--tag', self.args.worker_tag])
if self.args.worker_group:
command.extend(['--group', self.args.worker_group])
if self.args.worker_exit_after_num_runs and self.args.worker_exit_after_num_runs > 0:
command.extend([
'--exit-after-num-runs',
str(self.args.worker_exit_after_num_runs)
])
if self.args.worker_max_work_dir_size:
command.extend(
['--max-work-dir-size', self.args.worker_max_work_dir_size])
if self.args.worker_delete_work_dir_on_exit:
command.extend(['--delete-work-dir-on-exit'])
if self.args.worker_exit_on_exception:
command.extend(['--exit-on-exception'])
if self.args.worker_tag_exclusive:
command.extend(['--tag-exclusive'])
if self.args.worker_pass_down_termination:
command.extend(['--pass-down-termination'])
return command
def run_loop(self):
while True:
try:
self.run_one_iteration()
except (
urllib.error.URLError,
http.client.HTTPException,
socket.error,
NotFoundError,
JsonApiException,
):
# Sometimes, network errors occur when running the WorkerManager . These are often
# transient exceptions, and retrying the command would lead to success---as a result,
# we ignore these network-based exceptions (rather than fatally exiting from the
# WorkerManager )
traceback.print_exc()
if self.args.once:
break
logger.debug('Sleeping {} seconds'.format(self.args.sleep_time))
time.sleep(self.args.sleep_time)
def run_one_iteration(self):
# Get staged bundles for the current user. The principle here is that we want to get all of
# the staged bundles can be run by this user.
keywords = ['state=' + State.STAGED] + self.args.search
# If the current user is "codalab", don't filter by .mine because the workers owned
# by "codalab" can be shared by all users. But, for all other users, we only
# want to see their staged bundles.
if os.environ.get('CODALAB_USERNAME') != "codalab":
keywords += [".mine"]
# The keywords below search for `request_queue=<worker tag>` OR `request_queue=tag=<worker tag>`
# If support for this is removed so that 'request_queue' is always set to be '<worker tag>'
# (and not tag=<worker tag>) this search query can again be simplified.
# NOTE: server/bundle_manager.py has the server-side matching logic that should be synced
# with this search request.
if self.args.worker_tag_exclusive and self.args.worker_tag:
keywords += [
"request_queue=%s,tag=%s" %
(self.args.worker_tag, self.args.worker_tag)
]
bundles = self.codalab_client.fetch('bundles',
params={
'worksheet': None,
'keywords': keywords,
'include': ['owner']
})
new_staged_uuids = [bundle['uuid'] for bundle in bundles]
old_staged_uuids = self.staged_uuids
# Bundles that were staged but now aren't
removed_uuids = [
uuid for uuid in old_staged_uuids if uuid not in new_staged_uuids
]
self.staged_uuids = new_staged_uuids
logger.info('Staged bundles [{}]: {}'.format(
' '.join(keywords), ' '.join(self.staged_uuids) or '(none)'))
# Get worker jobs
worker_jobs = self.get_worker_jobs()
pending_worker_jobs, active_worker_jobs = [], []
for job in worker_jobs:
(active_worker_jobs
if job.active else pending_worker_jobs).append(job)
# Print status
logger.info(
'{} staged bundles ({} removed since last time), {} worker jobs (min={}, max={}) ({} active, {} pending)'
.format(
len(self.staged_uuids),
len(removed_uuids),
len(worker_jobs),
self.args.min_workers,
self.args.max_workers,
len(active_worker_jobs),
len(pending_worker_jobs),
))
want_workers = False
# There is a staged bundle AND there aren't any workers that are still booting up/starting
if len(self.staged_uuids) > 0:
logger.info(
'Want to launch a worker because we have {} > 0 staged bundles'
.format(len(self.staged_uuids)))
want_workers = True
if want_workers:
# Make sure we don't launch workers too quickly.
seconds_since_last_worker = int(time.time() -
self.last_worker_start_time)
if seconds_since_last_worker < self.args.min_seconds_between_workers:
logger.info(
'Don\'t launch because waited {} < {} seconds since last worker'
.format(seconds_since_last_worker,
self.args.min_seconds_between_workers))
want_workers = False
# Make sure we don't queue up more workers than staged UUIDs if there are
# more workers still booting up than staged bundles
if len(pending_worker_jobs) >= len(self.staged_uuids):
logger.info(
'Don\'t launch because still more pending workers than staged bundles ({} >= {})'
.format(len(pending_worker_jobs), len(self.staged_uuids)))
want_workers = False
# Don't launch more than `max_workers`.
# For now, only the number of workers is used to determine what workers
# we launch.
if len(worker_jobs) >= self.args.max_workers:
logger.info(
'Don\'t launch because too many workers already ({} >= {})'
.format(len(worker_jobs), self.args.max_workers))
want_workers = False
# We have fewer than min_workers, so launch one regardless of other constraints
if len(worker_jobs) < self.args.min_workers:
logger.info(
'Launch a worker because we are under the minimum ({} < {})'.
format(len(worker_jobs), self.args.min_workers))
want_workers = True
if want_workers:
logger.info('Starting a worker!')
self.start_worker_job()
self.last_worker_start_time = time.time()
def main(args):
manager = CodaLabManager()
model = manager.model()
# Get the the message
subject = args.subject
with open(args.body_file) as f:
body_template = f.read()
mime_type = 'html' if args.body_file.endswith('.html') else 'plain'
# Figure out who we want to send
to_send_list = get_to_send_list(model, args.threshold)
sent_list = get_sent_list(args.sent_file)
sent_emails = set(info['email'] for info in sent_list)
pending_to_send_list = [info for info in to_send_list if info['email'] not in sent_emails]
print('Already sent %d emails, %d to go' % (len(sent_list), len(pending_to_send_list)))
for i, info in enumerate(pending_to_send_list):
if args.only_email and args.only_email != info['email']:
continue
# Derived fields
info['greeting_name'] = info['first_name'] or info['last_name'] or info['user_name']
info['full_name'] = ' '.join([x for x in [info['first_name'], info['last_name']] if x])
info['email_description'] = (
'%s <%s>' % (info['full_name'], info['email']) if info['full_name'] else info['email']
)
info['sent_time'] = time.time()
print(
'Sending %s/%s (%s>=%s, doit=%s): [%s] %s'
% (
i,
len(pending_to_send_list),
info['notifications'],
args.threshold,
args.doit,
info['user_name'],
info['email_description'],
)
)
# Apply template to get body of message
body = body_template
for field, value in info.items():
body = body.replace('{{' + field + '}}', unicode(value or ''))
if args.verbose >= 1:
print('To : %s' % info['email_description'])
print('Subject : %s' % subject)
print(body)
print('-------')
if not args.doit:
continue
# Send the actual email
manager.emailer.send_email(
recipient=info['email_description'], subject=subject, body=body, mime_type=mime_type
)
# Record that we sent
with open(args.sent_file, 'a') as f:
print >> f, json.dumps(info)
f.flush()
def run_command(
args,
expected_exit_code=0,
max_output_chars=1024,
env=None,
include_stderr=False,
binary=False,
force_subprocess=False,
cwd=None,
):
# We import the following imports here because codalab_service.py imports TestModule from
# this file. If we kept the imports at the top, then anyone who ran codalab_service.py
# would also have to install all the dependencies that BundleCLI and CodaLabManager use.
from codalab.lib.bundle_cli import BundleCLI
from codalab.lib.codalab_manager import CodaLabManager
def sanitize(string, max_chars=256):
# Sanitize and truncate output so it can be printed on the command line.
# Don't print out binary.
if isinstance(string, bytes):
string = '<binary>'
if len(string) > max_chars:
string = string[:max_chars] + ' (...more...)'
return string
# If we don't care about the exit code, set `expected_exit_code` to None.
print(">>", *map(str, args), sep=" ")
sys.stdout.flush()
try:
kwargs = dict(env=env)
if not binary:
kwargs = dict(kwargs, encoding="utf-8")
if include_stderr:
kwargs = dict(kwargs, stderr=subprocess.STDOUT)
if cwd:
kwargs = dict(kwargs, cwd=cwd)
if not force_subprocess:
# In this case, run the Codalab CLI directly, which is much faster
# than opening a new subprocess to do so.
stderr = io.StringIO() # Not used; we just don't want to redirect cli.stderr to stdout.
stdout = FakeStdout()
cli = BundleCLI(CodaLabManager(), stdout=stdout, stderr=stderr)
try:
cli.do_command(args[1:])
exitcode = 0
except SystemExit as e:
exitcode = e.code
output = stdout.getvalue()
else:
output = subprocess.check_output([a.encode() for a in args], **kwargs)
exitcode = 0
except subprocess.CalledProcessError as e:
output = e.output
exitcode = e.returncode
except Exception:
output = traceback.format_exc()
exitcode = 'test-cli exception'
if expected_exit_code is not None and exitcode != expected_exit_code:
colorize = Colorizer.red
extra = ' BAD'
else:
colorize = Colorizer.cyan
extra = ''
print(colorize(" (exit code %s, expected %s%s)" % (exitcode, expected_exit_code, extra)))
sys.stdout.flush()
print(sanitize(output, max_output_chars))
sys.stdout.flush()
assert expected_exit_code == exitcode, 'Exit codes don\'t match'
return output.rstrip()
class WorkerManager(object):
"""
The abstract class for a worker manager. Different backends like AWS,
Azure, Slurm should override `get_worker_jobs` and `start_worker_job`.
The basic architecture of the WorkerManager is extremely simple: to a
first-order approximation, it simply launches `cl-worker`s as AWS/Azure
batch jobs as long as there are staged bundles.
The simplicity means that we don't need to manage state about how workers
and bundles are related - that logic is complex and is done by the usual
worker system, and we don't want to duplicate that for every possible
worker backend.
More specifically, a worker manager will monitor a job queue to see how
many worker jobs are running, and try to keep that between `min_workers`
and `max_workers`. It will also monitor the staged bundles that satisfy a
certain `search` criterion. If there are staged bundles then it will issue
a `start_worker_job()` call, provided some other conditions are met (e.g.,
don't start workers too fast).
The WorkerManager is all client-side code, so it can be customized as one
sees fit.
Notes:
- The worker manager is not visible via CodaLab (i.e., CodaLab has no
notion of a worker manager or what it's trying to do - all it sees is
bundles and workers). One needs to monitor the AWS/Azure Batch system
separately.
- Resource handling is not currently supported. Generally, the safe thing
is to create a separate queue for different resource needs and put the
burden of deciding on the user.
"""
# Subcommand name to use for this worker manager type
NAME = 'worker-manager'
DESCRIPTION = 'Base class for Worker Managers, please implement for your deployment'
@staticmethod
def add_arguments_to_subparser(subparser):
"""
Add any arguments specific to this worker manager to the given subparser
"""
raise NotImplementedError
def __init__(self, args):
self.args = args
self.codalab_manager = CodaLabManager()
self.codalab_client = self.codalab_manager.client(args.server)
self.staged_uuids = []
self.last_worker_start_time = 0
logger.info('Started worker manager.')
def get_worker_jobs(self):
"""Return a list of `WorkerJob`s."""
raise NotImplementedError
def start_worker_job(self):
"""Start a new `WorkerJob`."""
raise NotImplementedError
def run_loop(self):
while True:
self.run_one_iteration()
if self.args.once:
break
logger.debug('Sleeping {} seconds'.format(self.args.sleep_time))
time.sleep(self.args.sleep_time)
def run_one_iteration(self):
# Get staged bundles for the current user.
keywords = ['state=' + State.STAGED] + [".mine"] + self.args.search
if self.args.worker_tag:
keywords.append('request_queue=tag=' + self.args.worker_tag)
bundles = self.codalab_client.fetch('bundles',
params={
'worksheet': None,
'keywords': keywords,
'include': ['owner']
})
new_staged_uuids = [bundle['uuid'] for bundle in bundles]
old_staged_uuids = self.staged_uuids
# Bundles that were staged but now aren't
removed_uuids = [
uuid for uuid in old_staged_uuids if uuid not in new_staged_uuids
]
self.staged_uuids = new_staged_uuids
logger.info('Staged bundles [{}]: {}'.format(
' '.join(keywords), ' '.join(self.staged_uuids) or '(none)'))
# Get worker jobs
worker_jobs = self.get_worker_jobs()
pending_worker_jobs, active_worker_jobs = [], []
for job in worker_jobs:
(active_worker_jobs
if job.active else pending_worker_jobs).append(job)
# Print status
logger.info(
'{} staged bundles ({} removed since last time), {} worker jobs (min={}, max={}) ({} active, {} pending)'
.format(
len(self.staged_uuids),
len(removed_uuids),
len(worker_jobs),
self.args.min_workers,
self.args.max_workers,
len(active_worker_jobs),
len(pending_worker_jobs),
))
want_workers = False
# There is a staged bundle AND there aren't any workers that are still booting up/starting
if len(self.staged_uuids) > 0:
logger.info(
'Want to launch a worker because we have {} > 0 staged bundles'
.format(len(self.staged_uuids)))
want_workers = True
if want_workers:
# Make sure we don't launch workers too quickly.
seconds_since_last_worker = int(time.time() -
self.last_worker_start_time)
if seconds_since_last_worker < self.args.min_seconds_between_workers:
logger.info(
'Don\'t launch becaused waited {} < {} seconds since last worker'
.format(seconds_since_last_worker,
self.args.min_seconds_between_workers))
want_workers = False
# Make sure we don't queue up more workers than staged UUIDs if there are
# more workers still booting up than staged bundles
if len(pending_worker_jobs) >= len(self.staged_uuids):
logger.info(
'Don\'t launch because still more pending workers than staged bundles ({} >= {})'
.format(len(pending_worker_jobs), len(self.staged_uuids)))
want_workers = False
# Don't launch more than `max_workers`.
# For now, only the number of workers is used to determine what workers
# we launch.
if len(worker_jobs) >= self.args.max_workers:
logger.info(
'Don\'t launch because too many workers already ({} >= {})'
.format(len(worker_jobs), self.args.max_workers))
want_workers = False
# We have fewer than min_workers, so launch one regardless of other constraints
if len(worker_jobs) < self.args.min_workers:
logger.info(
'Launch a worker because we are under the minimum ({} < {})'.
format(len(worker_jobs), self.args.min_workers))
want_workers = True
if want_workers:
logger.info('Starting a worker!')
self.start_worker_job()
self.last_worker_start_time = time.time()
class DryRunAbort(Exception):
"""Raised at end of transaction of dry run."""
def __str__(self):
return """
This was a dry run, no migration occurred. To perform full migration,
run again with `-f':
%s -f
""".rstrip() % sys.argv[0]
dry_run = False if len(sys.argv) > 1 and sys.argv[1] == '-f' else True
manager = CodaLabManager()
model = manager.model()
CODALAB_HOME = manager.codalab_home
# Turn on query logging
model.engine.echo = True
###############################################################
# Configure connection to Django database
###############################################################
django_config = read_json_or_die(os.path.join(CODALAB_HOME,
'website-config.json'))
# Use default settings as defined in codalab-worksheets
if 'database' not in django_config:
#!./venv/bin/python
"""
Script that creates the root user.
"""
import sys
sys.path.append('.')
import getpass
from codalab.lib import crypt_util
from codalab.lib.codalab_manager import CodaLabManager
from codalab.objects.user import User
manager = CodaLabManager()
model = manager.model()
username = manager.root_user_name()
user_id = manager.root_user_id()
if len(sys.argv) == 2:
password = sys.argv[1]
else:
while True:
password = getpass.getpass()
if getpass.getpass('Config password: '******'Passwords don\'t match. Try again.'
print
def _create_cli(self, worksheet_uuid):
output_buffer = StringIO()
manager = CodaLabManager(temporary=True, clients={'local': self.client})
manager.set_current_worksheet_uuid(self.client, worksheet_uuid)
cli = bundle_cli.BundleCLI(manager, headless=True, stdout=output_buffer, stderr=output_buffer)
return cli, output_buffer
@Commands.command(
'bundle-manager',
help='Start the bundle manager that executes run and make bundles.',
arguments=(Commands.Argument(
'--sleep-time',
help='Number of seconds to wait between successive actions.',
type=int,
default=0.5), ),
)
def do_bundle_manager_command(bundle_cli, args):
bundle_cli._fail_if_headless(args)
from codalab.worker.bundle_manager import BundleManager
manager = BundleManager.create(bundle_cli.manager)
# Register a signal handler to ensure safe shutdown.
for sig in [signal.SIGTERM, signal.SIGINT, signal.SIGHUP]:
signal.signal(sig, lambda signup, frame: manager.signal())
manager.run(args.sleep_time)
if __name__ == '__main__':
cli = BundleCLI(CodaLabManager())
try:
cli.do_command(sys.argv[1:])
except KeyboardInterrupt:
print 'Terminated by Ctrl-C'
sys.exit(130)
#!./venv/bin/python
"""
Script that creates the default CodaLab OAuth2 clients.
- codalab_cli_client for the Bundle CLI clients authenticating through the Password Grant
- codalab_worker_client for workers authenticating through the Password Grant
TODO(skoo): Create row for the web client given a redirect url.
"""
import sys
sys.path.append('.')
from codalab.lib.codalab_manager import CodaLabManager
from codalab.objects.oauth2 import OAuth2Client
manager = CodaLabManager()
model = manager.model()
if not model.get_oauth2_client('codalab_cli_client'):
model.save_oauth2_client(OAuth2Client(
model,
client_id='codalab_cli_client',
secret=None,
name='Codalab CLI',
user_id=None,
grant_type='password',
response_type='token',
scopes='default',
redirect_uris='',
))
Indexes the contents of the bundle store. Used during the launch of the new
worker system.
TODO(klopyrev): Delete once it's launched.
"""
import sys
sys.path.append('.')
from codalab.common import State
from codalab.lib.codalab_manager import CodaLabManager
from codalab.model.tables import bundle as cl_bundle, bundle_contents_index as cl_bundle_contents_index
from sqlalchemy import distinct, select
from worker.file_util import index_contents
manager = CodaLabManager()
bundle_store = manager.bundle_store()
model = manager.model()
engine = model.engine
with engine.begin() as conn:
bundles = conn.execute(
select([cl_bundle.c.uuid])
.where(cl_bundle.c.state.in_([State.READY, State.FAILED]))
).fetchall()
indexed_bundles = conn.execute(
select([distinct(cl_bundle_contents_index.c.bundle_uuid)])
).fetchall()
uuids_to_index = (set(bundle.uuid for bundle in bundles) -
set(bundle.bundle_uuid for bundle in indexed_bundles))
class DryRunAbort(Exception):
"""Raised at end of transaction of dry run."""
def __str__(self):
return ("""
This was a dry run, no migration occurred. To perform full migration,
run again with `-f':
%s -f
""".rstrip() % sys.argv[0])
dry_run = False if len(sys.argv) > 1 and sys.argv[1] == '-f' else True
manager = CodaLabManager()
model = manager.model()
CODALAB_HOME = manager.codalab_home
# Turn on query logging
model.engine.echo = True
###############################################################
# Configure connection to Django database
###############################################################
django_config = read_json_or_die(
os.path.join(CODALAB_HOME, 'website-config.json'))
# Use default settings as defined in codalab-worksheets
if 'database' not in django_config:
django_config['database'] = {
class TestBase:
"""
Base class for BundleManager tests with a CodaLab Manager hitting an in-memory SQLite database.
"""
def setUp(self):
self.codalab_manager = CodaLabManager()
self.codalab_manager.config['server']['class'] = 'SQLiteModel'
self.bundle_manager = BundleManager(self.codalab_manager)
self.download_manager = self.codalab_manager.download_manager()
self.upload_manager = self.codalab_manager.upload_manager()
self.user_id = generate_uuid()
self.bundle_manager._model.add_user(
"codalab",
"*****@*****.**",
"Test",
"User",
"password",
"Stanford",
user_id=self.user_id,
)
def create_make_bundle(self, state=State.MAKING):
"""Creates a MakeBundle with the given state."""
bundle = MakeBundle.construct(
targets=[],
command='',
metadata=BASE_METADATA_MAKE_BUNDLE,
owner_id=self.user_id,
uuid=generate_uuid(),
state=state,
)
return bundle
def save_bundle(self, bundle):
"""Saves the given bundle to the database."""
self.bundle_manager._model.save_bundle(bundle)
def read_bundle(self, bundle, extra_path=""):
"""Retrieves the given bundle from the bundle store and returns
its contents.
Args:
extra_path: path appended to bundle store location from which to read the file.
Returns:
Bundle contents
"""
with open(
os.path.join(
self.codalab_manager.bundle_store().get_bundle_location(
bundle.uuid), extra_path),
"r",
) as f:
return f.read()
def write_bundle(self, bundle, contents=""):
"""Writes the given contents to the location of the given bundle.
Args:
bundle: bundle to write
contents: string to write
Returns:
None
"""
with open(
self.codalab_manager.bundle_store().get_bundle_location(
bundle.uuid), "w+") as f:
f.write(contents)
def update_bundle(self, bundle, update):
return self.bundle_manager._model.update_bundle(bundle, update)
def create_run_bundle(self, state=State.CREATED, metadata=None):
"""Creates a RunBundle.
Args:
state: state for the new bundle
metadata: additional metadata to add to the bundle.
"""
bundle = RunBundle.construct(
targets=[],
command='',
metadata=dict(BASE_METADATA, **(metadata or {})),
owner_id=self.user_id,
uuid=generate_uuid(),
state=state,
)
bundle.is_anonymous = False
return bundle
def create_bundle_single_dep(self,
parent_state=State.READY,
bundle_state=State.CREATED,
bundle_type=RunBundle):
"""Creates a bundle with a single dependency, which is mounted at path "src" of the
new bundle.
Args:
parent_state: State of the parent bundle. Defaults to State.READY.
bundle_state: State of the new bundle. Defaults to State.CREATED.
bundle_type: Type of child bundle to create; valid values are RunBundle and MakeBundle. Defaults to RunBundle.
Returns:
(bundle, parent)
"""
parent = self.create_run_bundle(parent_state)
self.write_bundle(parent, FILE_CONTENTS_1)
bundle = (self.create_run_bundle(bundle_state) if bundle_type
== RunBundle else self.create_make_bundle(bundle_state))
bundle.dependencies = [
Dependency({
"parent_uuid": parent.uuid,
"parent_path": "",
"child_uuid": bundle.uuid,
"child_path": "src",
})
]
return bundle, parent
def create_bundle_two_deps(self):
"""Create a bundle with two dependencies. The first dependency is mounted at path "src1"
and the second is mounted at path "src2" of the new bundle.
Returns:
(bundle, parent1, parent2)
"""
parent1 = self.create_run_bundle(state=State.READY)
self.write_bundle(parent1, FILE_CONTENTS_1)
parent2 = self.create_run_bundle(state=State.READY)
self.write_bundle(parent2, FILE_CONTENTS_2)
bundle = MakeBundle.construct(
targets=[],
command='',
metadata=BASE_METADATA_MAKE_BUNDLE,
owner_id=self.user_id,
uuid=generate_uuid(),
state=State.STAGED,
)
bundle.dependencies = [
Dependency({
"parent_uuid": parent1.uuid,
"parent_path": "",
"child_uuid": bundle.uuid,
"child_path": "src1",
}),
Dependency({
"parent_uuid": parent2.uuid,
"parent_path": "",
"child_uuid": bundle.uuid,
"child_path": "src2",
}),
]
return bundle, parent1, parent2
def mock_worker_checkin(self,
cpus=0,
gpus=0,
memory_bytes=0,
free_disk_bytes=0,
tag=None,
user_id=None):
"""Perform a mock check-in of a new worker."""
worker_id = generate_uuid()
self.bundle_manager._worker_model.worker_checkin(
user_id=user_id
or self.bundle_manager._model.root_user_id, # codalab-owned worker
worker_id=worker_id,
tag=tag,
group_name=None,
cpus=cpus,
gpus=gpus,
memory_bytes=memory_bytes,
free_disk_bytes=free_disk_bytes,
dependencies=[],
shared_file_system=False,
tag_exclusive=False,
exit_after_num_runs=999999999,
is_terminating=False,
)
# Mock a reply from the worker
self.bundle_manager._worker_model.send_json_message = Mock(
return_value=True)
return worker_id
def mock_bundle_checkin(self, bundle, worker_id, user_id=None):
"""Mock a worker checking in with the latest state of a bundle.
Args:
bundle: Bundle to check in.
worker_id ([type]): worker id of the worker that performs the checkin.
user_id (optional): user id that performs the checkin. Defaults to the default user id.
"""
worker_run = BundleCheckinState(
uuid=bundle.uuid,
run_status="",
bundle_start_time=0,
container_time_total=0,
container_time_user=0,
container_time_system=0,
docker_image="",
state=bundle.state,
remote="",
exitcode=0,
failure_message="",
cpu_usage=0.0,
memory_limit=0,
)
self.bundle_manager._model.bundle_checkin(bundle, worker_run, user_id
or self.user_id, worker_id)
#!./venv/bin/python
"""
Script that creates the default CodaLab OAuth2 clients.
- codalab_cli_client for the Bundle CLI clients authenticating through the Password Grant
- codalab_worker_client for workers authenticating through the Password Grant
TODO(skoo): Create row for the web client given a redirect url.
"""
import sys
sys.path.append('.')
from codalab.lib.codalab_manager import CodaLabManager
from codalab.objects.oauth2 import OAuth2Client
manager = CodaLabManager()
model = manager.model()
if not model.get_oauth2_client('codalab_cli_client'):
model.save_oauth2_client(
OAuth2Client(
model,
client_id='codalab_cli_client',
secret=None,
name='Codalab CLI',
user_id=None,
grant_type='password',
response_type='token',
scopes='default',
redirect_uris='',
))
#!./venv/bin/python
"""
Script that creates the root user.
"""
import sys
sys.path.append('.')
import getpass
from codalab.lib import crypt_util
from codalab.lib.codalab_manager import CodaLabManager
from codalab.objects.user import User
manager = CodaLabManager()
model = manager.model()
username = manager.root_user_name()
user_id = manager.root_user_id()
if len(sys.argv) == 2:
password = sys.argv[1]
else:
while True:
password = getpass.getpass('Password for %s(%s): ' % (username, user_id))
if getpass.getpass('Confirm password: '******'Passwords don\'t match. Try again.')
if model.get_user(user_id=user_id, check_active=False):
update = {
def _create_cli(self, worksheet_uuid):
manager = CodaLabManager(temporary=True, clients={settings.BUNDLE_SERVICE_URL: self.client})
manager.set_current_worksheet_uuid(self.client, worksheet_uuid)
cli = bundle_cli.BundleCLI(manager, headless=True)
return cli
def main(args):
manager = CodaLabManager()
# Get the the message
subject = args.subject
with open(args.body_file) as f:
body_template = f.read()
if args.body_file.endswith('.md'):
body_template = f"""
<div style='margin:auto; width: 100%; max-width: 600px'>
<img src=https://worksheets.codalab.org/img/codalab-logo.png style='max-width: 100%;' />
<h1>CodaLab Worksheets</h1>
{markdown2.markdown(body_template)}<br><br>
<small>If you'd like stop receiving these emails, please <a href='https://worksheets.codalab.org/account/profile'>update your account settings on CodaLab</a>.</small>
</div>
"""
mime_type = ('html' if args.body_file.endswith('.html')
or args.body_file.endswith('.md') else 'plain')
# Figure out who we want to send
to_send_list = ([{
'email': e,
'first_name': '',
'last_name': '',
'user_name': '',
'notifications': 2
} for e in args.emails.split(",")] if args.emails else get_to_send_list(
manager.model(), args.threshold))
sent_list = get_sent_list(args.sent_file)
sent_emails = set(info['email'] for info in sent_list)
pending_to_send_list = [
info for info in to_send_list if info['email'] not in sent_emails
]
print('Already sent %d emails, %d to go' %
(len(sent_list), len(pending_to_send_list)))
for i, info in enumerate(pending_to_send_list):
if args.only_email and args.only_email != info['email']:
continue
# Derived fields
info['greeting_name'] = info['first_name'] or info[
'last_name'] or info['user_name']
info['full_name'] = ' '.join(
[x for x in [info['first_name'], info['last_name']] if x])
info['email_description'] = ('%s <%s>' %
(info['full_name'], info['email'])
if info['full_name'] else info['email'])
info['sent_time'] = time.time()
print(('Sending %s/%s (%s>=%s, doit=%s): [%s] %s' % (
i,
len(pending_to_send_list),
info['notifications'],
args.threshold,
args.doit,
info['user_name'],
info['email_description'],
)))
# Apply template to get body of message
body = body_template
for field, value in info.items():
body = body.replace('{{' + field + '}}', str(value or ''))
if args.verbose >= 1:
print('To : %s' % info['email_description'])
print('Subject : %s' % subject)
print(body)
print('-------')
if not args.doit:
continue
# Send the actual email
manager.emailer.send_email(recipient=info['email_description'],
subject=subject,
body=body,
mime_type=mime_type)
# Record that we sent
with open(args.sent_file, 'a') as f:
print(json.dumps(info), file=f)
f.flush()
