def _get_container_pod_env_vars(self, persistence_outputs,
persistence_data, outputs_refs_jobs,
outputs_refs_experiments, ephemeral_token):
logs_path = self._get_logs_path()
outputs_path = self._get_outputs_path(
persistence_outputs=persistence_outputs)
env_vars = get_job_env_vars(
namespace=self.namespace,
persistence_outputs=persistence_outputs,
outputs_path=outputs_path,
persistence_data=persistence_data,
log_level=self.log_level,
logs_path=logs_path,
outputs_refs_jobs=outputs_refs_jobs,
outputs_refs_experiments=outputs_refs_experiments,
ephemeral_token=ephemeral_token,
)
return env_vars + [
get_env_var(name=constants.CONFIG_MAP_CLUSTER_KEY_NAME,
value=json.dumps(self.cluster_def)),
get_env_var(name=constants.CONFIG_MAP_PARAMS_KEY_NAME,
value=self.params),
get_env_var(name=constants.CONFIG_MAP_EXPERIMENT_INFO_KEY_NAME,
value=json.dumps(self.experiment_labels)),
]
def test_env_vars(self):
# String value
env_var = get_env_var(name='foo', value='bar')
assert env_var.name == 'foo'
assert env_var.value == 'bar'
# Int value
env_var = get_env_var(name='foo', value=1)
assert env_var.name == 'foo'
assert env_var.value == '1'
# Dict value
env_var = get_env_var(name='foo', value={'moo': 'bar'})
assert env_var.name == 'foo'
assert env_var.value == '{"moo": "bar"}'
def get_env_vars(self, task_type, task_idx):
if task_type == TaskType.MASTER:
rank = 0
master_addr = 'localhost'
else:
rank = task_idx + 1
master_addr = self.resource_manager.get_resource_name(task_type=TaskType.MASTER,
task_idx=0)
env_vars = [
get_env_var(name='MASTER_ADDR', value=master_addr),
get_env_var(name='MASTER_PORT', value=self.ports[0]),
get_env_var(name='WORLD_SIZE', value=self.get_n_pods(TaskType.WORKER) + 1),
get_env_var(name='RANK', value=rank)
]
return env_vars
def get_sidecar_container(job_container_name,
sidecar_container_name,
sidecar_docker_image,
sidecar_docker_image_pull_policy,
namespace,
sidecar_config,
sidecar_args,
internal_health_check_url,
internal_reconcile_url,
volume_mounts,
env_vars=None):
"""Return a pod sidecar container."""
env_vars = to_list(env_vars) if env_vars else []
env_vars += get_sidecar_env_vars(
namespace=namespace,
job_container_name=job_container_name,
internal_health_check_url=internal_health_check_url,
internal_reconcile_url=internal_reconcile_url)
for k, v in sidecar_config.items():
env_vars.append(get_env_var(name=k, value=v))
return client.V1Container(
name=sidecar_container_name,
image=sidecar_docker_image,
image_pull_policy=sidecar_docker_image_pull_policy,
command=get_sidecar_command(),
env=env_vars,
volume_mounts=volume_mounts,
resources=get_sidecar_resources(),
args=sidecar_args)
def get_sidecar_env_vars(namespace, job_container_name,
internal_health_check_url, internal_reconcile_url):
env_vars = get_internal_env_vars(
namespace=namespace,
service_internal_header=InternalServices.SIDECAR,
authentication_type=AuthenticationTypes.INTERNAL_TOKEN,
include_internal_token=True)
return env_vars + [
get_env_var(name='POLYAXON_CONTAINER_ID', value=job_container_name),
get_env_var(name='POLYAXON_INTERNAL_HEALTH_CHECK_URL',
value=internal_health_check_url),
get_env_var(name='POLYAXON_INTERNAL_RECONCILE_URL',
value=internal_reconcile_url),
get_from_field_ref(name='POLYAXON_POD_ID', field_path='metadata.name'),
]
def _get_container_pod_env_vars(self, persistence_outputs,
persistence_data, outputs_refs_jobs,
outputs_refs_experiments, ephemeral_token):
return [
get_env_var(name=constants.CONFIG_MAP_JOB_INFO_KEY_NAME,
value=json.dumps(self.labels)),
]
def get_env_vars(self, task_type, task_idx):
tf_config = {
'cluster': self.get_cluster(),
'task': {'type': task_type, 'index': task_idx},
'model_dir': stores.get_experiment_outputs_path(
persistence=self.persistence_config.outputs,
experiment_name=self.experiment_name,
cloning_strategy=self.cloning_strategy),
'environment': 'cloud'
}
return get_env_var(name='TF_CONFIG', value=tf_config)
def get_init_container(self, init_command, init_args, env_vars,
context_mounts, persistence_outputs,
persistence_data):
"""Pod init container for setting outputs path."""
env_vars = to_list(env_vars, check_none=True)
env_vars += [
get_env_var(name=constants.CONFIG_MAP_JOB_INFO_KEY_NAME,
value=json.dumps(self.labels)),
]
return client.V1Container(
name=self.init_container_name,
image=self.init_docker_image,
image_pull_policy=self.init_docker_image_pull_policy,
command=init_command,
env=env_vars,
args=init_args,
resources=get_init_resources(),
volume_mounts=context_mounts)
def get_env_vars(self, task_type, task_idx):
role = TaskType.SCHEDULER if task_type == TaskType.MASTER else task_type
env_vars = [
get_env_var(name='DMLC_NUM_WORKER',
value=self.get_n_pods(TaskType.WORKER)),
get_env_var(name='DMLC_NUM_SERVER',
value=self.get_n_pods(TaskType.SERVER)),
get_env_var(name='DMLC_PS_ROOT_URI',
value=self.resource_manager.get_resource_name(
task_type=TaskType.MASTER, task_idx=0)),
get_env_var(name='DMLC_PS_ROOT_PORT', value=self.ports[0]),
get_env_var(name='DMLC_ROLE', value=role)
]
if task_type == TaskType.SERVER:
env_vars.append(get_env_var(name='DMLC_SERVER_ID', value=task_idx))
elif task_type == TaskType.WORKER:
env_vars.append(get_env_var(name='DMLC_SERVER_ID', value=0))
env_vars.append(get_env_var(name='DMLC_WORKER_ID', value=task_idx))
return env_vars
def _get_container_pod_env_vars(self, persistence_outputs,
persistence_data, outputs_refs_jobs,
outputs_refs_experiments, ephemeral_token):
logs_path = self._get_logs_path()
outputs_path = self._get_outputs_path(
persistence_outputs=persistence_outputs)
env_vars = get_job_env_vars(
namespace=self.namespace,
persistence_outputs=persistence_outputs,
outputs_path=outputs_path,
persistence_data=persistence_data,
log_level=self.log_level,
logs_path=logs_path,
outputs_refs_jobs=outputs_refs_jobs,
outputs_refs_experiments=outputs_refs_experiments,
ephemeral_token=ephemeral_token,
)
return env_vars + [
get_env_var(name=constants.CONFIG_MAP_NOTEBOOK_INFO_KEY_NAME,
value=json.dumps(self.labels)),
]
