def create_env(self, taskdef: TaskDefinition) -> dict:
"""
Create a container environment dict from a task definition.
Arguments:
taskdef (TaskDefinition): Task definition
Returns:
env (dict): Environment variable dict
"""
env = {
**taskdef.env,
ENV_GZIP_ENABLED: '1',
ENV_TASK_CLUSTER: env_pack(self.serialize()),
ENV_TASK_DEFINITION: env_pack(taskdef.serialize()),
}
# check total length of environment data
length = 0
for key, value in env.items():
length += len(str(key)) + len(str(value))
if length > MAX_ENV_LENGTH:
raise ProviderError(f'Task environment too long. Was {length}, max: {MAX_ENV_LENGTH}')
return env
async def send_init(self, taskdef: TaskDefinition) -> None:
"""
Send a task initialization message.
Arguments:
taskdef (TaskDefinition): New task definition
"""
await self.msg(TASK_INIT, task=taskdef.serialize())
def test_create_docker_task():
dp = DockerProvider()
docker = dp.docker
env_vars = {'hello': 'team'}
taskdef = TaskDefinition(
name=TEST_TASK,
image=TEST_IMAGE,
parent='parent',
env=env_vars,
inputs={
'hello': '123',
'child': False,
},
# disables any output.
# this is hacky and should be refactored
# we need a proper way to disable all logging
upstream='disabled',
)
# run task
task = dp.spawn(taskdef)
assert task.id == taskdef.id
assert hasattr(task, 'container')
assert hasattr(task.container, 'id')
# try to grab the container from docker api
container = docker.containers.get(task.container.id)
assert task.container == container
# make sure container is properly labeled
assert container.labels == {
LABEL_TASK_ID: task.id,
LABEL_PARENT_ID: 'parent',
}
# wait for container to execute
result = container.wait()
assert result['StatusCode'] == 0
# test task will dump info as json, so we can pick it up
# make sure it matches what we put in.
logs = container.logs()
task_dump = json.loads(logs)
# taskdef
assert taskdef.serialize() == task_dump['taskdef']
# actual environment variables
for key, val in env_vars.items():
assert task_dump['env'][key] == val
def test_create_docker_task():
dp = DockerProvider()
docker = dp.docker
env_vars = {'hello': 'team'}
taskdef = TaskDefinition(
name=TEST_TASK,
image=TEST_IMAGE,
parent='parent',
env=env_vars,
inputs={
'hello': '123',
'child': False,
},
)
# run task
task = dp.spawn(taskdef)
assert task.id == taskdef.id
assert hasattr(task, 'container')
assert hasattr(task.container, 'id')
# try to grab the container from docker api
container = docker.containers.get(task.container.id)
assert task.container == container
# make sure container is properly labeled
assert container.labels[LABEL_TASK_ID] == task.id
assert container.labels[LABEL_PARENT_ID] == 'parent'
# test task will dump info as json, so we can pick it up
# make sure it matches what we put in.
task_dump = None
for msg in task.logs():
print(msg)
if msg['type'] == 'task/log':
task_dump = json.loads(msg['data'])
# wait for container to execute
result = container.wait()
assert result['StatusCode'] == 0
# taskdef
assert task_dump is not None
assert taskdef.serialize() == task_dump['taskdef']
# actual environment variables
for key, val in env_vars.items():
assert task_dump['env'][key] == val
def create_env(self, taskdef: TaskDefinition) -> dict:
"""
Create a container environment dict from a task definition.
Arguments:
taskdef (TaskDefinition): Task definition
Returns:
env (dict): Environment variable dict
"""
return {
**taskdef.env,
ENV_TASK_CLUSTER: json.dumps(self.serialize()),
ENV_TASK_DEFINITION: json.dumps(taskdef.serialize()),
}
def base_environment(cluster, taskdef: TaskDefinition) -> dict:
"""
Create a container environment dict from a task definition.
Arguments:
taskdef (TaskDefinition): Task definition
Returns:
env (dict): Environment variable dict
"""
return {
**taskdef.env,
ENV_GZIP_ENABLED: '1',
ENV_TASK_CLUSTER: env_pack(cluster.serialize()),
ENV_TASK_DEFINITION: env_pack(taskdef.serialize()),
}
def spawn(self, taskdef: TaskDefinition) -> RemoteTask:
try:
task = self.rpc('spawn', **taskdef.serialize())
return RemoteTask(TaskDefinition.deserialize(task), self)
except RpcError as e:
raise TaskCreationError(str(e))
