async def setup(
self,
runtime_env: "RuntimeEnv", # noqa: F821
context: RuntimeEnvContext,
logger: Optional[logging.Logger] = default_logger,
):
if not runtime_env.has_py_container(
) or not runtime_env.py_container_image():
return
container_driver = "podman"
container_command = [
container_driver,
"run",
"-v",
self._ray_tmp_dir + ":" + self._ray_tmp_dir,
"--cgroup-manager=cgroupfs",
"--network=host",
"--pid=host",
"--ipc=host",
"--env-host",
]
container_command.append("--env")
container_command.append("RAY_RAYLET_PID=" +
os.getenv("RAY_RAYLET_PID"))
if runtime_env.py_container_run_options():
container_command.extend(runtime_env.py_container_run_options())
# TODO(chenk008): add resource limit
container_command.append("--entrypoint")
container_command.append("python")
container_command.append(runtime_env.py_container_image())
context.py_executable = " ".join(container_command)
logger.info("start worker in container with prefix: {}".format(
context.py_executable))
def setup(self,
runtime_env: RuntimeEnv,
context: RuntimeEnvContext,
logger: Optional[logging.Logger] = default_logger):
if not runtime_env.has_conda() and not runtime_env.has_pip():
return
logger.debug("Setting up conda or pip for runtime_env: "
f"{runtime_env.serialize()}")
if runtime_env.conda_env_name():
conda_env_name = runtime_env.conda_env_name()
else:
conda_dict = get_conda_dict(runtime_env, self._resources_dir)
protocol, hash = parse_uri(runtime_env.conda_uri())
conda_env_name = self._get_path_from_hash(hash)
assert conda_dict is not None
ray_pip = current_ray_pip_specifier(logger=logger)
if ray_pip:
extra_pip_dependencies = [ray_pip, "ray[default]"]
elif runtime_env.get_extension("_inject_current_ray") == "True":
extra_pip_dependencies = (
_resolve_install_from_source_ray_dependencies())
else:
extra_pip_dependencies = []
conda_dict = inject_dependencies(conda_dict, _current_py_version(),
extra_pip_dependencies)
# It is not safe for multiple processes to install conda envs
# concurrently, even if the envs are different, so use a global
# lock for all conda installs.
# See https://github.com/ray-project/ray/issues/17086
file_lock_name = "ray-conda-install.lock"
with FileLock(os.path.join(self._resources_dir, file_lock_name)):
try:
conda_yaml_file = os.path.join(self._resources_dir,
"environment.yml")
with open(conda_yaml_file, "w") as file:
yaml.dump(conda_dict, file)
if conda_env_name in self._created_envs:
logger.debug(f"Conda env {conda_env_name} already "
"created, skipping creation.")
else:
create_conda_env(conda_yaml_file,
prefix=conda_env_name,
logger=logger)
self._created_envs.add(conda_env_name)
finally:
os.remove(conda_yaml_file)
if runtime_env.get_extension("_inject_current_ray"):
_inject_ray_to_conda_site(conda_path=conda_env_name,
logger=logger)
context.py_executable = "python"
context.command_prefix += get_conda_activate_commands(conda_env_name)
logger.info(
f"Finished setting up runtime environment at {conda_env_name}")
def modify_context(
self,
uris: List[str],
runtime_env: "RuntimeEnv", # noqa: F821
context: RuntimeEnvContext,
logger: Optional[logging.Logger] = default_logger,
):
if not runtime_env.has_pip():
return
# PipPlugin only uses a single URI.
uri = uris[0]
# Update py_executable.
protocol, hash = parse_uri(uri)
target_dir = self._get_path_from_hash(hash)
virtualenv_python = _PathHelper.get_virtualenv_python(target_dir)
if not os.path.exists(virtualenv_python):
raise ValueError(
f"Local directory {target_dir} for URI {uri} does "
"not exist on the cluster. Something may have gone wrong while "
"installing the runtime_env `pip` packages.")
context.py_executable = virtualenv_python
context.command_prefix += [
_PathHelper.get_virtualenv_activate_command(target_dir)
]
def start_specific_server(self, client_id: str,
job_config: JobConfig) -> bool:
"""
Start up a RayClient Server for an incoming client to
communicate with. Returns whether creation was successful.
"""
specific_server = self._get_server_for_client(client_id)
assert specific_server, f"Server has not been created for: {client_id}"
output, error = self.node.get_log_file_handles(
f"ray_client_server_{specific_server.port}", unique=True)
serialized_runtime_env = job_config.get_serialized_runtime_env()
runtime_env_config = job_config.get_proto_runtime_env_config()
if not serialized_runtime_env or serialized_runtime_env == "{}":
# TODO(edoakes): can we just remove this case and always send it
# to the agent?
serialized_runtime_env_context = RuntimeEnvContext().serialize()
else:
serialized_runtime_env_context = self._create_runtime_env(
serialized_runtime_env=serialized_runtime_env,
runtime_env_config=runtime_env_config,
specific_server=specific_server,
)
proc = start_ray_client_server(
self.address,
self.node.node_ip_address,
specific_server.port,
stdout_file=output,
stderr_file=error,
fate_share=self.fate_share,
server_type="specific-server",
serialized_runtime_env_context=serialized_runtime_env_context,
redis_password=self._redis_password,
)
# Wait for the process being run transitions from the shim process
# to the actual RayClient Server.
pid = proc.process.pid
if sys.platform != "win32":
psutil_proc = psutil.Process(pid)
else:
psutil_proc = None
# Don't use `psutil` on Win32
while psutil_proc is not None:
if proc.process.poll() is not None:
logger.error(
f"SpecificServer startup failed for client: {client_id}")
break
cmd = psutil_proc.cmdline()
if _match_running_client_server(cmd):
break
logger.debug(
"Waiting for Process to reach the actual client server.")
time.sleep(0.5)
specific_server.set_result(proc)
logger.info(f"SpecificServer started on port: {specific_server.port} "
f"with PID: {pid} for client: {client_id}")
return proc.process.poll() is None
def run_setup_with_logger():
runtime_env: dict = json.loads(serialized_runtime_env or "{}")
allocated_resource: dict = json.loads(
serialized_allocated_resource_instances or "{}")
# Use a separate logger for each job.
per_job_logger = self.get_or_create_logger(request.job_id)
# TODO(chenk008): Add log about allocated_resource to
# avoid lint error. That will be moved to cgroup plugin.
per_job_logger.debug(f"Worker has resource :"
f"{allocated_resource}")
context = RuntimeEnvContext(
env_vars=runtime_env.get("env_vars"))
self._conda_manager.setup(
runtime_env, context, logger=per_job_logger)
self._working_dir_manager.setup(
runtime_env, context, logger=per_job_logger)
# Add the mapping of URIs -> the serialized environment to be
# used for cache invalidation.
for uri in runtime_env.get("uris") or []:
self._working_dir_uri_to_envs[uri].add(
serialized_runtime_env)
# Run setup function from all the plugins
for plugin_class_path in runtime_env.get("plugins", {}).keys():
plugin_class = import_attr(plugin_class_path)
# TODO(simon): implement uri support
plugin_class.create("uri not implemented", runtime_env,
context)
plugin_class.modify_context("uri not implemented",
runtime_env, context)
return context
def setup(self,
runtime_env: RuntimeEnv,
context: RuntimeEnvContext,
logger: Optional[logging.Logger] = default_logger):
if not runtime_env.has_pip():
return
logger.debug(f"Setting up pip for runtime_env: {runtime_env}")
pip_packages: List[str] = runtime_env.pip_packages()
target_dir = self._get_path_from_hash(_get_pip_hash(pip_packages))
_install_pip_list_to_dir(pip_packages, target_dir, logger=logger)
# Despite Ray being removed from the input pip list during validation,
# other packages in the pip list (for example, xgboost_ray) may
# themselves include Ray as a dependency. In this case, we will have
# inadvertently installed the latest Ray version in the target_dir,
# which may cause Ray version mismatch issues. Uninstall it here, if it
# exists, to make the workers use the Ray that is already
# installed in the cluster.
#
# In the case where the user explicitly wants to include Ray in their
# pip list (and signals this by setting the environment variable below)
# then we don't want this deletion logic, so we skip it.
if os.environ.get(RAY_RUNTIME_ENV_ALLOW_RAY_IN_PIP) != 1:
ray_path = Path(target_dir) / "ray"
if ray_path.exists() and ray_path.is_dir():
shutil.rmtree(ray_path)
# Insert the target directory into the PYTHONPATH.
python_path = target_dir
if "PYTHONPATH" in context.env_vars:
python_path += os.pathsep + context.env_vars["PYTHONPATH"]
context.env_vars["PYTHONPATH"] = python_path
def modify_context(uri: str, plugin_config_dict: dict,
ctx: RuntimeEnvContext) -> None:
ctx.env_vars[MyPlugin.env_key] = str(plugin_config_dict["env_value"])
ctx.command_prefix.append(
f"echo {plugin_config_dict['tmp_content']} > "
f"{plugin_config_dict['tmp_file']}")
ctx.py_executable = (plugin_config_dict["prefix_command"] + " " +
ctx.py_executable)
def test_inherit_cluster_env_pythonpath(monkeypatch):
monkeypatch.setenv("PYTHONPATH",
"last" + os.pathsep + os.environ.get("PYTHONPATH", ""))
context = RuntimeEnvContext(env_vars={"PYTHONPATH": "middle"})
set_pythonpath_in_context("first", context)
assert context.env_vars["PYTHONPATH"].startswith(
os.pathsep.join(["first", "middle", "last"]))
def run_setup_with_logger():
runtime_env = RuntimeEnv(
serialized_runtime_env=serialized_runtime_env)
allocated_resource: dict = json.loads(
serialized_allocated_resource_instances or "{}")
# Use a separate logger for each job.
per_job_logger = self.get_or_create_logger(request.job_id)
# TODO(chenk008): Add log about allocated_resource to
# avoid lint error. That will be moved to cgroup plugin.
per_job_logger.debug(f"Worker has resource :"
f"{allocated_resource}")
context = RuntimeEnvContext(env_vars=runtime_env.env_vars())
self._pip_manager.setup(
runtime_env, context, logger=per_job_logger)
self._conda_manager.setup(
runtime_env, context, logger=per_job_logger)
self._py_modules_manager.setup(
runtime_env, context, logger=per_job_logger)
self._working_dir_manager.setup(
runtime_env, context, logger=per_job_logger)
self._container_manager.setup(
runtime_env, context, logger=per_job_logger)
# Add the mapping of URIs -> the serialized environment to be
# used for cache invalidation.
if runtime_env.working_dir_uri():
uri = runtime_env.working_dir_uri()
self._uris_to_envs[uri].add(serialized_runtime_env)
if runtime_env.py_modules_uris():
for uri in runtime_env.py_modules_uris():
self._uris_to_envs[uri].add(serialized_runtime_env)
if runtime_env.conda_uri():
uri = runtime_env.conda_uri()
self._uris_to_envs[uri].add(serialized_runtime_env)
if runtime_env.pip_uri():
uri = runtime_env.pip_uri()
self._uris_to_envs[uri].add(serialized_runtime_env)
if runtime_env.plugin_uris():
for uri in runtime_env.plugin_uris():
self._uris_to_envs[uri].add(serialized_runtime_env)
# Run setup function from all the plugins
for plugin_class_path, config in runtime_env.plugins():
logger.debug(
f"Setting up runtime env plugin {plugin_class_path}")
plugin_class = import_attr(plugin_class_path)
# TODO(simon): implement uri support
plugin_class.create("uri not implemented",
json.loads(config), context)
plugin_class.modify_context("uri not implemented",
json.loads(config), context)
return context
def modify_context(
self,
uris: List[str],
plugin_config_dict: dict,
ctx: RuntimeEnvContext,
logger: logging.Logger,
) -> None:
if PRIORITY_TEST_ENV_VAR_NAME in ctx.env_vars:
raise RuntimeError(
f"Env var {PRIORITY_TEST_ENV_VAR_NAME} has been set to "
f"{ctx.env_vars[PRIORITY_TEST_ENV_VAR_NAME]}.")
ctx.env_vars[
PRIORITY_TEST_ENV_VAR_NAME] = PriorityTestPlugin2.env_value
def modify_context(
self,
uris: List[str],
plugin_config_dict: dict,
ctx: RuntimeEnvContext,
logger: logging.Logger,
) -> None:
if PRIORITY_TEST_ENV_VAR_NAME in ctx.env_vars:
ctx.env_vars[
PRIORITY_TEST_ENV_VAR_NAME] += PriorityTestPlugin1.env_value
else:
ctx.env_vars[
PRIORITY_TEST_ENV_VAR_NAME] = PriorityTestPlugin1.env_value
def modify_context(
self,
uris: List[str],
plugin_config_dict: dict,
ctx: RuntimeEnvContext,
logger: logging.Logger,
) -> None:
ctx.env_vars[MyPlugin.env_key] = str(plugin_config_dict["env_value"])
ctx.command_prefix.append(
f"echo {plugin_config_dict['tmp_content']} > "
f"{plugin_config_dict['tmp_file']}")
ctx.py_executable = (plugin_config_dict["prefix_command"] + " " +
ctx.py_executable)
def setup(self,
runtime_env: dict,
context: RuntimeEnvContext,
logger: Optional[logging.Logger] = default_logger):
if not runtime_env.get("conda") and not runtime_env.get("pip"):
return
logger.debug(f"Setting up conda or pip for runtime_env: {runtime_env}")
conda_dict = get_conda_dict(runtime_env, self._resources_dir)
if isinstance(runtime_env.get("conda"), str):
conda_env_name = runtime_env["conda"]
else:
assert conda_dict is not None
ray_pip = current_ray_pip_specifier(logger=logger)
if ray_pip:
extra_pip_dependencies = [ray_pip, "ray[default]"]
elif runtime_env.get("_inject_current_ray"):
extra_pip_dependencies = (
_resolve_install_from_source_ray_dependencies())
else:
extra_pip_dependencies = []
conda_dict = inject_dependencies(conda_dict, _current_py_version(),
extra_pip_dependencies)
logger.info(f"Setting up conda environment with {runtime_env}")
# It is not safe for multiple processes to install conda envs
# concurrently, even if the envs are different, so use a global
# lock for all conda installs.
# See https://github.com/ray-project/ray/issues/17086
file_lock_name = "ray-conda-install.lock"
with FileLock(os.path.join(self._resources_dir, file_lock_name)):
conda_dir = os.path.join(self._resources_dir, "conda")
try_to_create_directory(conda_dir)
conda_yaml_path = os.path.join(conda_dir, "environment.yml")
with open(conda_yaml_path, "w") as file:
# Sort keys because we hash based on the file contents,
# and we don't want the hash to depend on the order
# of the dependencies.
yaml.dump(conda_dict, file, sort_keys=True)
conda_env_name = get_or_create_conda_env(conda_yaml_path,
conda_dir,
logger=logger)
if runtime_env.get("_inject_current_ray"):
conda_path = os.path.join(conda_dir, conda_env_name)
_inject_ray_to_conda_site(conda_path, logger=logger)
context.py_executable = "python"
context.command_prefix += get_conda_activate_commands(conda_env_name)
logger.info(
f"Finished setting up runtime environment at {conda_env_name}")
def set_pythonpath_in_context(python_path: str, context: RuntimeEnvContext):
"""Insert the path as the first entry in PYTHONPATH in the runtime env.
This is compatible with users providing their own PYTHONPATH in env_vars,
and is also compatible with the existing PYTHONPATH in the cluster.
The import priority is as follows:
this python_path arg > env_vars PYTHONPATH > existing cluster env PYTHONPATH.
"""
if "PYTHONPATH" in context.env_vars:
python_path += os.pathsep + context.env_vars["PYTHONPATH"]
if "PYTHONPATH" in os.environ:
python_path += os.pathsep + os.environ["PYTHONPATH"]
context.env_vars["PYTHONPATH"] = python_path
async def _setup_runtime_env(runtime_env, serialized_runtime_env,
serialized_allocated_resource_instances):
allocated_resource: dict = json.loads(
serialized_allocated_resource_instances or "{}")
# Use a separate logger for each job.
per_job_logger = self.get_or_create_logger(request.job_id)
# TODO(chenk008): Add log about allocated_resource to
# avoid lint error. That will be moved to cgroup plugin.
per_job_logger.debug(f"Worker has resource :"
f"{allocated_resource}")
context = RuntimeEnvContext(env_vars=runtime_env.env_vars())
await self._container_manager.setup(runtime_env,
context,
logger=per_job_logger)
for manager in self._base_plugin_cache_managers.values():
await manager.create_if_needed(runtime_env,
context,
logger=per_job_logger)
def setup_plugins():
# Run setup function from all the plugins
if runtime_env.plugins():
for (
setup_context
) in self._runtime_env_plugin_manager.sorted_plugin_setup_contexts(
runtime_env.plugins()):
per_job_logger.debug(
f"Setting up runtime env plugin {setup_context.name}"
)
# TODO(architkulkarni): implement uri support
setup_context.class_instance.validate(runtime_env)
setup_context.class_instance.create(
"uri not implemented", setup_context.config,
context)
setup_context.class_instance.modify_context(
"uri not implemented",
setup_context.config,
context,
per_job_logger,
)
loop = asyncio.get_event_loop()
# Plugins setup method is sync process, running in other threads
# is to avoid blocking asyncio loop
await loop.run_in_executor(None, setup_plugins)
return context
def modify_context(
self,
uri: str,
runtime_env: "RuntimeEnv", # noqa: F821
context: RuntimeEnvContext,
logger: Optional[logging.Logger] = default_logger,
):
if not runtime_env.has_conda():
return
if runtime_env.conda_env_name():
conda_env_name = runtime_env.conda_env_name()
else:
protocol, hash = parse_uri(runtime_env.conda_uri())
conda_env_name = self._get_path_from_hash(hash)
context.py_executable = "python"
context.command_prefix += get_conda_activate_commands(conda_env_name)
def setup(self,
runtime_env: dict,
context: RuntimeEnvContext,
logger: Optional[logging.Logger] = default_logger):
if not runtime_env.get("working_dir"):
return
working_dir = download_and_unpack_package(runtime_env["working_dir"],
self._resources_dir,
logger=logger)
context.command_prefix += [f"cd {working_dir}"]
# Insert the working_dir as the first entry in PYTHONPATH. This is
# compatible with users providing their own PYTHONPATH in env_vars.
python_path = working_dir
if "PYTHONPATH" in context.env_vars:
python_path += os.pathsep + context.env_vars["PYTHONPATH"]
context.env_vars["PYTHONPATH"] = python_path
def setup(self,
runtime_env: RuntimeEnv,
context: RuntimeEnvContext,
logger: Optional[logging.Logger] = default_logger):
if not runtime_env.py_modules():
return
module_dirs = []
for uri in runtime_env.py_modules():
module_dir = download_and_unpack_package(
uri, self._resources_dir, logger=logger)
module_dirs.append(module_dir)
# Insert the py_modules directories into the PYTHONPATH.
python_path = os.pathsep.join(module_dirs)
if "PYTHONPATH" in context.env_vars:
python_path += os.pathsep + context.env_vars["PYTHONPATH"]
context.env_vars["PYTHONPATH"] = python_path
def setup(self,
runtime_env: dict,
context: RuntimeEnvContext,
logger: Optional[logging.Logger] = default_logger):
if not runtime_env.get("uris"):
return
working_dir = self.ensure_runtime_env_setup(runtime_env["uris"],
logger=logger)
if working_dir is None:
return
context.command_prefix += [f"cd {working_dir}"]
# Insert the working_dir as the first entry in PYTHONPATH. This is
# compatible with users providing their own PYTHONPATH in env_vars.
python_path = working_dir
if "PYTHONPATH" in context.env_vars:
python_path += os.pathsep + context.env_vars["PYTHONPATH"]
context.env_vars["PYTHONPATH"] = python_path
def modify_context(self, uri: Optional[str], runtime_env_dict: Dict,
context: RuntimeEnvContext):
if uri is None:
return
local_dir = get_local_dir_from_uri(uri, self._resources_dir)
if not local_dir.exists():
raise ValueError(
f"Local directory {local_dir} for URI {uri} does "
"not exist on the cluster. Something may have gone wrong while "
"downloading or unpacking the working_dir.")
context.command_prefix += [f"cd {local_dir}"]
# Insert the working_dir as the first entry in PYTHONPATH. This is
# compatible with users providing their own PYTHONPATH in env_vars.
python_path = str(local_dir)
if "PYTHONPATH" in context.env_vars:
python_path += os.pathsep + context.env_vars["PYTHONPATH"]
context.env_vars["PYTHONPATH"] = python_path
def modify_context(
self,
uri: str,
runtime_env: RuntimeEnv,
context: RuntimeEnvContext,
logger: Optional[logging.Logger] = default_logger,
):
if not runtime_env.has_pip():
return
# Update py_executable.
protocol, hash = parse_uri(uri)
target_dir = self._get_path_from_hash(hash)
virtualenv_python = _PathHelper.get_virtualenv_python(target_dir)
if not os.path.exists(virtualenv_python):
raise ValueError(
f"Local directory {target_dir} for URI {uri} does "
"not exist on the cluster. Something may have gone wrong while "
"installing the runtime_env `pip` packages.")
context.py_executable = virtualenv_python
def modify_context(
self,
uri: str,
runtime_env: RuntimeEnv,
context: RuntimeEnvContext,
logger: Optional[logging.Logger] = default_logger,
):
if not runtime_env.has_pip():
return
# Insert the target directory into the PYTHONPATH.
protocol, hash = parse_uri(uri)
target_dir = get_local_dir_from_uri(uri, self._resources_dir)
if not target_dir.exists():
raise ValueError(
f"Local directory {target_dir} for URI {uri} does "
"not exist on the cluster. Something may have gone wrong while "
"installing the runtime_env `pip` packages.")
python_path = str(target_dir)
if "PYTHONPATH" in context.env_vars:
python_path += os.pathsep + context.env_vars["PYTHONPATH"]
context.env_vars["PYTHONPATH"] = python_path
async def test_create_delete_size_equal(tmpdir, ray_start_regular):
"""Tests that `create` and `delete_uri` return the same size for a URI."""
# Create an arbitrary nonempty directory to upload.
path = Path(tmpdir)
dir_to_upload = path / "dir_to_upload"
dir_to_upload.mkdir(parents=True)
filepath = dir_to_upload / "file"
with filepath.open("w") as file:
file.write("F" * 100)
uri = get_uri_for_directory(dir_to_upload)
assert get_directory_size_bytes(dir_to_upload) > 0
uploaded = upload_package_if_needed(uri, tmpdir, dir_to_upload)
assert uploaded
manager = WorkingDirManager(tmpdir)
created_size_bytes = await manager.create(uri, {}, RuntimeEnvContext())
deleted_size_bytes = manager.delete_uri(uri)
assert created_size_bytes == deleted_size_bytes
def modify_context(
self,
uris: Optional[List[str]],
runtime_env_dict: Dict,
context: RuntimeEnvContext,
logger: Optional[logging.Logger] = default_logger,
):
if uris is None:
return
module_dirs = []
for uri in uris:
module_dir = get_local_dir_from_uri(uri, self._resources_dir)
if not module_dir.exists():
raise ValueError(
f"Local directory {module_dir} for URI {uri} does "
"not exist on the cluster. Something may have gone wrong while "
"downloading or unpacking the py_modules files.")
module_dirs.append(str(module_dir))
# Insert the py_modules directories into the PYTHONPATH.
python_path = os.pathsep.join(module_dirs)
if "PYTHONPATH" in context.env_vars:
python_path += os.pathsep + context.env_vars["PYTHONPATH"]
context.env_vars["PYTHONPATH"] = python_path
from ray.core.generated.common_pb2 import Language
logger = logging.getLogger(__name__)
parser = argparse.ArgumentParser(description=(
"Set up the environment for a Ray worker and launch the worker."))
parser.add_argument("--serialized-runtime-env",
type=str,
help="the serialized parsed runtime env dict")
parser.add_argument("--serialized-runtime-env-context",
type=str,
help="the serialized runtime env context")
parser.add_argument("--language",
type=str,
help="the language type of the worker")
if __name__ == "__main__":
args, remaining_args = parser.parse_known_args()
# NOTE(chenk008): we still need it to start worker in container.
remaining_args.append("--worker-shim-pid={}".format(os.getpid()))
# NOTE(edoakes): args.serialized_runtime_env_context is only None when
# we're starting the main Ray client proxy server. That case should
# probably not even go through this codepath.
runtime_env_context = RuntimeEnvContext.deserialize(
args.serialized_runtime_env_context or "{}")
runtime_env_context.exec_worker(remaining_args,
Language.Value(args.language))
def run_setup_with_logger():
runtime_env = RuntimeEnv(serialized_runtime_env=serialized_runtime_env)
allocated_resource: dict = json.loads(
serialized_allocated_resource_instances or "{}"
)
# Use a separate logger for each job.
per_job_logger = self.get_or_create_logger(request.job_id)
# TODO(chenk008): Add log about allocated_resource to
# avoid lint error. That will be moved to cgroup plugin.
per_job_logger.debug(f"Worker has resource :" f"{allocated_resource}")
context = RuntimeEnvContext(env_vars=runtime_env.env_vars())
self._container_manager.setup(
runtime_env, context, logger=per_job_logger
)
for (manager, uri_cache) in [
(self._working_dir_manager, self._working_dir_uri_cache),
(self._conda_manager, self._conda_uri_cache),
(self._pip_manager, self._pip_uri_cache),
]:
uri = manager.get_uri(runtime_env)
if uri is not None:
if uri not in uri_cache:
per_job_logger.debug(f"Cache miss for URI {uri}.")
size_bytes = manager.create(
uri, runtime_env, context, logger=per_job_logger
)
uri_cache.add(uri, size_bytes, logger=per_job_logger)
else:
per_job_logger.debug(f"Cache hit for URI {uri}.")
uri_cache.mark_used(uri, logger=per_job_logger)
manager.modify_context(uri, runtime_env, context)
# Set up py_modules. For now, py_modules uses multiple URIs so
# the logic is slightly different from working_dir, conda, and
# pip above.
py_modules_uris = self._py_modules_manager.get_uris(runtime_env)
if py_modules_uris is not None:
for uri in py_modules_uris:
if uri not in self._py_modules_uri_cache:
per_job_logger.debug(f"Cache miss for URI {uri}.")
size_bytes = self._py_modules_manager.create(
uri, runtime_env, context, logger=per_job_logger
)
self._py_modules_uri_cache.add(
uri, size_bytes, logger=per_job_logger
)
else:
per_job_logger.debug(f"Cache hit for URI {uri}.")
self._py_modules_uri_cache.mark_used(
uri, logger=per_job_logger
)
self._py_modules_manager.modify_context(
py_modules_uris, runtime_env, context
)
# Add the mapping of URIs -> the serialized environment to be
# used for cache invalidation.
if runtime_env.working_dir_uri():
uri = runtime_env.working_dir_uri()
self._uris_to_envs[uri].add(serialized_runtime_env)
if runtime_env.py_modules_uris():
for uri in runtime_env.py_modules_uris():
self._uris_to_envs[uri].add(serialized_runtime_env)
if runtime_env.conda_uri():
uri = runtime_env.conda_uri()
self._uris_to_envs[uri].add(serialized_runtime_env)
if runtime_env.pip_uri():
uri = runtime_env.pip_uri()
self._uris_to_envs[uri].add(serialized_runtime_env)
if runtime_env.plugin_uris():
for uri in runtime_env.plugin_uris():
self._uris_to_envs[uri].add(serialized_runtime_env)
# Run setup function from all the plugins
for plugin_class_path, config in runtime_env.plugins():
per_job_logger.debug(
f"Setting up runtime env plugin {plugin_class_path}"
)
plugin_class = import_attr(plugin_class_path)
# TODO(simon): implement uri support
plugin_class.create(
"uri not implemented", json.loads(config), context
)
plugin_class.modify_context(
"uri not implemented", json.loads(config), context
)
return context
async def _setup_runtime_env(runtime_env, serialized_runtime_env,
serialized_allocated_resource_instances):
allocated_resource: dict = json.loads(
serialized_allocated_resource_instances or "{}")
# Use a separate logger for each job.
per_job_logger = self.get_or_create_logger(request.job_id)
# TODO(chenk008): Add log about allocated_resource to
# avoid lint error. That will be moved to cgroup plugin.
per_job_logger.debug(f"Worker has resource :"
f"{allocated_resource}")
context = RuntimeEnvContext(env_vars=runtime_env.env_vars())
await self._container_manager.setup(runtime_env,
context,
logger=per_job_logger)
for (manager, uri_cache) in [
(self._working_dir_manager, self._working_dir_uri_cache),
(self._conda_manager, self._conda_uri_cache),
(self._pip_manager, self._pip_uri_cache),
]:
uri = manager.get_uri(runtime_env)
if uri is not None:
if uri not in uri_cache:
per_job_logger.debug(f"Cache miss for URI {uri}.")
size_bytes = await manager.create(
uri, runtime_env, context, logger=per_job_logger)
uri_cache.add(uri, size_bytes, logger=per_job_logger)
else:
per_job_logger.debug(f"Cache hit for URI {uri}.")
uri_cache.mark_used(uri, logger=per_job_logger)
manager.modify_context(uri, runtime_env, context)
# Set up py_modules. For now, py_modules uses multiple URIs so
# the logic is slightly different from working_dir, conda, and
# pip above.
for (manager, uri_cache) in [
(self._java_jars_manager, self._java_jars_uri_cache),
(self._py_modules_manager, self._py_modules_uri_cache),
]:
uris = manager.get_uris(runtime_env)
if uris is not None:
per_job_logger.debug(f"URIs is not None, URI {uris}.")
for uri in uris:
if uri not in uri_cache:
per_job_logger.debug(f"Cache miss for URI {uri}.")
size_bytes = await manager.create(
uri,
runtime_env,
context,
logger=per_job_logger)
uri_cache.add(uri,
size_bytes,
logger=per_job_logger)
else:
per_job_logger.debug(f"Cache hit for URI {uri}.")
uri_cache.mark_used(uri, logger=per_job_logger)
manager.modify_context(uris, runtime_env, context)
def setup_plugins():
# Run setup function from all the plugins
for plugin_class_path, config in runtime_env.plugins():
per_job_logger.debug(
f"Setting up runtime env plugin {plugin_class_path}")
plugin_class = import_attr(plugin_class_path)
# TODO(simon): implement uri support
plugin_class.create("uri not implemented",
json.loads(config), context)
plugin_class.modify_context("uri not implemented",
json.loads(config), context)
loop = asyncio.get_event_loop()
# Plugins setup method is sync process, running in other threads
# is to avoid blocks asyncio loop
await loop.run_in_executor(None, setup_plugins)
return context