From d07da15ecedf311966d99000cd58731d8f08d7bd Mon Sep 17 00:00:00 2001 From: JosepSampe Date: Sat, 29 Aug 2026 14:37:00 +0200 Subject: [PATCH 1/4] Add a full unit test suite for the entire project --- CHANGELOG.md | 65 + docs/source/compute_config/localhost.md | 100 ++ lithops/config.py | 295 +++-- lithops/constants.py | 95 +- lithops/executors.py | 1090 +++++++++------ lithops/future.py | 477 ++++--- lithops/invokers.py | 700 ++++++---- lithops/job/__init__.py | 3 +- lithops/job/job.py | 639 ++++++--- lithops/job/partitioner.py | 622 ++++----- lithops/job/serialize.py | 434 +++--- lithops/localhost/__init__.py | 5 +- lithops/localhost/config.py | 81 +- lithops/localhost/utils.py | 188 +++ lithops/localhost/v1/localhost.py | 484 ++++--- lithops/localhost/v1/runner.py | 104 +- lithops/localhost/v2/localhost.py | 559 ++++---- lithops/localhost/v2/runner.py | 99 +- lithops/monitor.py | 430 ++++-- lithops/plots.py | 243 ++-- lithops/retries.py | 229 ++-- lithops/scripts/cleaner.py | 392 ++++-- lithops/scripts/cli.py | 1134 +++++++++------- .../backends/aws_batch/aws_batch.py | 2 +- .../backends/aws_lambda/aws_lambda.py | 45 +- .../azure_containers/azure_containers.py | 113 +- .../azure_functions/azure_functions.py | 51 +- .../backends/code_engine/code_engine.py | 2 +- .../backends/gcp_cloudrun/cloudrun.py | 4 +- .../backends/gcp_functions/gcp_functions.py | 4 +- lithops/serverless/backends/k8s/k8s.py | 8 +- .../serverless/backends/knative/knative.py | 15 +- lithops/serverless/serverless.py | 140 +- lithops/standalone/__init__.py | 2 +- lithops/standalone/keeper.py | 120 +- lithops/standalone/master.py | 541 +++++--- lithops/standalone/runner.py | 63 +- lithops/standalone/standalone.py | 518 +++++--- lithops/standalone/utils.py | 194 ++- lithops/standalone/worker.py | 318 +++-- lithops/storage/cloud_proxy.py | 199 ++- lithops/storage/storage.py | 489 ++++--- lithops/storage/utils.py | 147 +- lithops/tests/conftest.py | 17 + lithops/tests/functions.py | 19 + lithops/tests/test_config.py | 553 ++++++++ lithops/tests/test_constants.py | 117 ++ lithops/tests/test_executors.py | 596 +++++++++ lithops/tests/test_future.py | 392 ++++++ lithops/tests/test_invokers.py | 972 ++++++++++++++ lithops/tests/test_job.py | 1178 +++++++++++++++++ lithops/tests/test_joblib.py | 234 ++++ lithops/tests/test_localhost.py | 1002 ++++++++++++++ lithops/tests/test_map.py | 32 + lithops/tests/test_monitor.py | 661 +++++++++ lithops/tests/test_plots.py | 136 ++ lithops/tests/test_retries.py | 193 +++ lithops/tests/test_scripts.py | 574 ++++++++ lithops/tests/test_serverless.py | 122 ++ lithops/tests/test_standalone.py | 664 ++++++++++ lithops/tests/test_standalone_master.py | 824 ++++++++++++ lithops/tests/test_storage_layer.py | 430 ++++++ lithops/tests/test_util.py | 418 ++++++ lithops/tests/test_utils.py | 658 +++++++++ lithops/tests/test_wait.py | 399 ++++++ lithops/tests/test_worker.py | 1169 ++++++++++++++++ lithops/util/ibm_token_manager.py | 125 +- lithops/util/joblib/__init__.py | 10 +- lithops/util/joblib/lithops_backend.py | 291 ++-- lithops/util/metrics.py | 59 +- lithops/util/ssh_client.py | 210 ++- lithops/utils.py | 652 +++++---- lithops/wait.py | 481 ++++--- lithops/worker/handler.py | 446 +++++-- lithops/worker/invoker.py | 45 +- lithops/worker/jobrunner.py | 469 ++++--- lithops/worker/status.py | 157 ++- lithops/worker/utils.py | 234 ++-- setup.py | 3 +- 79 files changed, 20501 insertions(+), 5484 deletions(-) create mode 100644 lithops/localhost/utils.py create mode 100644 lithops/tests/test_config.py create mode 100644 lithops/tests/test_constants.py create mode 100644 lithops/tests/test_executors.py create mode 100644 lithops/tests/test_invokers.py create mode 100644 lithops/tests/test_job.py create mode 100644 lithops/tests/test_joblib.py create mode 100644 lithops/tests/test_localhost.py create mode 100644 lithops/tests/test_monitor.py create mode 100644 lithops/tests/test_plots.py create mode 100644 lithops/tests/test_scripts.py create mode 100644 lithops/tests/test_serverless.py create mode 100644 lithops/tests/test_standalone.py create mode 100644 lithops/tests/test_standalone_master.py create mode 100644 lithops/tests/test_storage_layer.py create mode 100644 lithops/tests/test_util.py create mode 100644 lithops/tests/test_utils.py create mode 100644 lithops/tests/test_wait.py create mode 100644 lithops/tests/test_worker.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 3c74f8372..7d9ea8fdd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,70 @@ # Changelog +## [v3.7.1.dev0] + +### Added + +- [Tests] Added a unit test suite for all non-backend modules (18 files, 876 tests). +- [Core] Added a `log_prefix()` helper for uniform log prefixes across core and backends. +- [Core] Added a cache of serialized functions to avoid re-uploading the same function. +- [Core] Added `ShutdownSafeStreamHandler` to avoid tracebacks when logging on a closed stream. +- [Localhost] Added `localhost/utils.py` with helpers shared by the v1 and v2 backends. +- [AWS Batch] Added `instance_types` config option for EC2/SPOT compute environments. + +### Changed + +- [Worker] Replaced the `multiprocessing` Manager queue of the worker pool with a POSIX pipe. +- [Core] Results under 8KB now travel in the call status instead of a separate storage object. +- [Core] Reorganised all non-backend modules for readability, with no behaviour change. +- [Core] `wait()` now returns two empty lists for empty input instead of `None`. +- [Core] `verify_args()` now raises a single message instead of a tuple. +- [Monitoring] The RabbitMQ queues of a call status now travel with the job. +- [CLI] `job list`, `worker list`, `image delete` and `image list` now reject unknown flags. +- [CLI] `lithops clean --all` no longer shadows the `all` builtin. +- [CLI] `lithops clean` now empties the local temp directory instead of removing it, and leaves the pending cleaner requests of the other processes alone. +- [Storage] `CloudFileProxy.walk()` now yields nothing for a missing path, like `os.walk`. +- [Storage] `cloud_open()` now raises `ValueError` on an unsupported mode. +- [Joblib] Capped the shared-argument upload and download pools at 32 threads. +- [Joblib] `lithops_args` is now applied to the pool that runs the batches. +- [Standalone] `docker login` now reads the password from stdin and quotes its arguments. + +### Fixed + +- [Localhost] Fixed a deadlock on a `map` after `wait()` and `get_result()`, caused by stale work queue sentinels. +- [Localhost] Fixed a partial `clear()` tearing down the consumers, tasks and latches of other jobs. +- [Localhost] Fixed a task starting after `stop()`, leaving a process nobody kills. +- [Localhost] Fixed the v2 job manager spinning a core while an invocation was queueing. +- [Localhost] Fixed two concurrent `invoke()` calls clearing each other's in-progress flag. +- [Localhost] Fixed the v2 container being removed while other jobs were still running in it. +- [Standalone] Fixed a dict race that killed the budget keeper and left the VM running. +- [Standalone] Fixed a file descriptor leak of the runner log, one per task. +- [Standalone] Fixed the worker `/stop` endpoint iterating the process map while it changed. +- [Standalone] Fixed `cancel_job_process()` raising on an emptied queue or a job with no queue. +- [Standalone] Fixed the master dropping the errors of its parallel worker and job requests. +- [Standalone] Fixed the SSH client keeping a client that failed to connect. +- [Storage] Fixed `delete_cloudobjects()` deleting part of the list before rejecting a foreign object. +- [Storage] Fixed `CloudFileProxy.listdir()` returning nothing for its default argument. +- [Core] Fixed `find_free_port()` setting `SO_REUSEADDR` after the bind. +- [Core] Fixed module inspection crashing on a function whose `__module__` is `None`. +- [Core] Fixed a hand-built `FuturesList` raising `AttributeError` instead of creating its executor. +- [Core] Fixed the cleaner skipping requests and two cleaners racing for the pid file. +- [Core] Fixed `lithops clean` deleting the local temp directory of the jobs running at the same time on the same machine. +- [Core] Fixed the cleaner reading a request another process was still writing. +- [Core] Fixed the cleaner looping forever on a request it could not read or classify. +- [Core] Fixed the cleaner lock surviving a killed cleaner and blocking every later one. +- [Monitoring] Fixed a nested executor publishing statuses to a queue nobody declares. +- [Monitoring] Fixed the failed RabbitMQ publishes being dropped with nothing in the log. +- [Worker] Fixed the memory monitor reporting a peak of zero where usage cannot be read. +- [Worker] Fixed the remote invoker returning before its invocations in flight were done. +- [Job] Fixed folder markers being counted as objects, returning empty partitions. +- [Joblib] Fixed the backend being unused with joblib 1.4+, which renamed `apply_async` to `submit`. +- [Joblib] Fixed a `KeyError` on shared arguments over 32KB, from a check-then-read on the disk cache. +- [Joblib] Fixed shared arguments going to the default storage instead of the configured one. +- [Joblib] Fixed a race losing one of two shared arguments proxied in the same call. +- [Joblib] Fixed `lithops[joblib]` missing `redis`, needed to import the backend. +- [IBM] Fixed the COS token manager raising if `ibm_botocore` hides the private expiry attribute. +- [Tests] Fixed the test suite depending on the order its files run in. + ## [v3.7.0] ### Added diff --git a/docs/source/compute_config/localhost.md b/docs/source/compute_config/localhost.md index 105a5f917..f2b3e2f9c 100644 --- a/docs/source/compute_config/localhost.md +++ b/docs/source/compute_config/localhost.md @@ -61,6 +61,106 @@ fexec = lithops.LocalhostExecutor(runtime='docker.io/lithopscloud/ibmcf-python-v In this mode of execution, you can use any Docker image that contains all the required dependencies. For example, the IBM Cloud Functions and Knative runtimes are compatible with it. +## Implementation versions (v1 and v2) + +There are two localhost implementations. **v2 is the default** (`localhost.version: 2`). Set `version: 1` only if you need the older job-at-a-time runner. + +```yaml +localhost: + version: 2 # default; use 1 for the alternative implementation +``` + +Both versions copy the Lithops package into `/tmp/lithops-/` and can use either the **default** Python interpreter or a **container** image. They differ in how they schedule activations. + +### How v2 works (default) + +v2 splits a job into **one task per function activation**. A pool of `worker_processes` consumer threads (default: CPU count) pulls tasks from an in-memory work queue and runs them in parallel. + +- **Default environment:** each task is a subprocess: `python localhost-runner.py run_job .task`. +- **Container environment:** Lithops starts **one** long-lived container (`docker run --detach` with `/bin/bash`) and runs each task with `docker exec … python3 … run_job`. The host `/tmp` tree is bind-mounted into the container so job files and the Lithops package are shared. + +### How v1 works + +v1 treats a Lithops **job** as a single unit. The client writes one JSON job file with all call IDs. A job-manager thread runs jobs **one after another** and waits for each process to exit. + +- **Default environment:** one subprocess runs the whole job: `python localhost-runner.py run_job .json`. Parallelism is inside that process (`multiprocessing`, `worker_processes` workers). +- **Container environment:** each job starts a **new** container (`docker run --name lithops_`). The container exits when the job finishes (`--rm`). There is no shared long-lived worker container. + +### v1 vs v2 + +| | **v2 (default)** | **v1** | +|---|---|---| +| Scheduling unit | One activation (call) | One job (all calls together) | +| Parallelism | `worker_processes` consumer threads, each running a task | Job manager is serial; parallelism is inside the job process | +| Default runtime | One Python subprocess per call | One Python subprocess per job | +| Container runtime | One detached container for the executor; `docker exec` per call | New `docker run` per job | +| Job payload on disk | Per-call `.task` files under `/tmp/lithops-*/jobs/` | One `.json` job file under the storage prefix | +| When to use | Default; better overlap of independent activations | Compatibility with the older runner | + +## Architecture diagram + +Localhost never provisions cloud VMs. The client, job manager, workers, and (optional) Docker engine all run on **your machine**. Function data uses localhost storage under `/tmp/lithops-/` unless you set `storage` to a remote backend. + +### v2 (default) + +```mermaid +flowchart TB + LAPTOP["Your laptop / FunctionExecutor"] + subgraph host [This machine] + H["LocalhostHandler v2"] + Q["Work queue\none JSON task per call"] + C1["Consumer thread 1"] + C2["Consumer thread N\nworker_processes"] + TMP["/tmp/lithops-user\npackage + jobs + logs"] + PY["python localhost-runner.py\nrun_job call.task"] + subgraph docker [Optional: one long-lived container] + CTR["docker run --detach --rm\nimage + /tmp mount"] + EXEC["docker exec python3\nrun_job call.task"] + end + end + STORAGE[(Localhost storage\nor S3 / COS / …)] + LAPTOP --> H + H -->|split calls| Q + Q --> C1 + Q --> C2 + C1 --> PY + C2 --> PY + C1 --> EXEC + C2 --> EXEC + PY --> TMP + EXEC --> CTR + CTR --> TMP + PY -->|read/write| STORAGE + EXEC -->|read/write| STORAGE +``` + +### v1 + +```mermaid +flowchart TB + LAPTOP["Your laptop / FunctionExecutor"] + subgraph host [This machine] + H["LocalhostHandler v1"] + JQ["Job queue\none JSON file per job"] + JM["Job manager thread\none job at a time"] + TMP["/tmp/lithops-user\npackage + job JSON + logs"] + PY["python localhost-runner.py\nrun_job job.json\nmultiprocessing workers"] + subgraph docker [Optional: new container per job] + RUN["docker run --name lithops_job\nimage + /tmp mount"] + end + end + STORAGE[(Localhost storage\nor S3 / COS / …)] + LAPTOP --> H + H -->|enqueue job file| JQ + JQ --> JM + JM --> PY + JM --> RUN + PY --> TMP + RUN --> TMP + PY -->|read/write| STORAGE + RUN -->|read/write| STORAGE +``` + ## Summary of configuration keys for Localhost: |Group|Key|Default|Mandatory|Additional info| diff --git a/lithops/config.py b/lithops/config.py index 522bdbf0a..68d7605be 100644 --- a/lithops/config.py +++ b/lithops/config.py @@ -24,17 +24,24 @@ from lithops import constants as c from lithops.version import __version__ from lithops.utils import CURRENT_PY_VERSION, get_mode, get_default_backend -from builtins import FileNotFoundError logger = logging.getLogger(__name__) -os.makedirs(c.LITHOPS_TEMP_DIR, exist_ok=True) -os.makedirs(c.JOBS_DIR, exist_ok=True) -os.makedirs(c.LOGS_DIR, exist_ok=True) -os.makedirs(c.CLEANER_DIR, exist_ok=True) +for _lithops_dir in (c.LITHOPS_TEMP_DIR, c.JOBS_DIR, c.LOGS_DIR, c.CLEANER_DIR): + os.makedirs(_lithops_dir, exist_ok=True) + +_USER_AGENT = f'lithops/{__version__}' +_LOCALHOST_FALLBACK = { + 'lithops': { + 'mode': c.LOCALHOST, + 'backend': c.LOCALHOST, + 'storage': c.LOCALHOST, + } +} def load_yaml_config(config_filename): + """Reads a YAML config file, or returns nothing if it does not exist""" import yaml try: with open(config_filename, 'r') as config_file: @@ -46,9 +53,11 @@ def load_yaml_config(config_filename): def dump_yaml_config(config_filename, data): + """Writes a config to a YAML file, creating its directory if needed""" import yaml - if not os.path.exists(os.path.dirname(config_filename)): - os.makedirs(os.path.dirname(config_filename)) + dirname = os.path.dirname(config_filename) + if dirname: + os.makedirs(dirname, exist_ok=True) with open(config_filename, "w") as config_file: yaml.dump(data, config_file, default_flow_style=False) @@ -56,29 +65,32 @@ def dump_yaml_config(config_filename, data): def get_default_config_filename(): """ - First checks .lithops_config - then checks LITHOPS_CONFIG_FILE environment variable - then ~/.lithops/config - and as last resort the global configuration /etc/lithops/config + Resolve the default Lithops config file, in order: + 1. LITHOPS_CONFIG_FILE environment variable + 2. .lithops_config in the current working directory + 3. ~/.lithops/config + 4. /etc/lithops/config """ if 'LITHOPS_CONFIG_FILE' in os.environ: - config_filename = os.environ['LITHOPS_CONFIG_FILE'] + return os.environ['LITHOPS_CONFIG_FILE'] - elif os.path.exists(".lithops_config"): - config_filename = os.path.abspath('.lithops_config') + if os.path.exists(".lithops_config"): + return os.path.abspath('.lithops_config') - else: - config_filename = c.CONFIG_FILE - if not os.path.exists(config_filename): - config_filename = c.CONFIG_FILE_GLOBAL - if not os.path.exists(config_filename): - return None + if os.path.exists(c.CONFIG_FILE): + return c.CONFIG_FILE + + if os.path.exists(c.CONFIG_FILE_GLOBAL): + return c.CONFIG_FILE_GLOBAL - return config_filename + return None def load_config(config_file=None, log=True): - """ Load the configuration """ + """ + Loads the configuration from a file, from the environment, or from the + default locations. Falls back to localhost mode when there is none + """ config_data = None if config_file: @@ -86,13 +98,15 @@ def load_config(config_file=None, log=True): if log: logger.debug(f"Loading configuration from {config_filename}") if not os.path.exists(config_filename): - raise FileNotFoundError(f"Config file {config_filename} doesn't exist") + raise FileNotFoundError( + f"Config file {config_filename} doesn't exist" + ) config_data = load_yaml_config(config_filename) elif 'LITHOPS_CONFIG' in os.environ: if log: logger.debug("Loading configuration from env LITHOPS_CONFIG") - config_data = json.loads(os.environ.get('LITHOPS_CONFIG')) + config_data = json.loads(os.environ['LITHOPS_CONFIG']) else: config_filename = get_default_config_filename() @@ -101,163 +115,210 @@ def load_config(config_file=None, log=True): logger.debug(f"Loading configuration from {config_filename}") config_data = load_yaml_config(config_filename) - if not config_data: # Set Lithops to Localhost mode + if not config_data: + # None, {}, or empty YAML all mean "no usable config" → localhost. + # A file containing `lithops: {}` is truthy and must NOT take this path. if log: - logger.debug("Config file not found. Setting Lithops to Localhost mode") - config_data = {'lithops': {'mode': c.LOCALHOST, 'backend': c.LOCALHOST, 'storage': c.LOCALHOST}} + logger.debug( + "Config file not found. Setting Lithops to Localhost mode" + ) + config_data = copy.deepcopy(_LOCALHOST_FALLBACK) return config_data -def get_log_info(config_file=None, config_data=None): - """ Return lithops logging information set in configuration """ - config_data = copy.deepcopy(config_data) or load_config(config_file, log=False) +def _copy_or_load_config(config_file, config_data, **load_kwargs): + # Treat None *and* {} as "no config provided" so callers fall back to + # load_config(). This is intentional: an empty dict must not skip file/env + # discovery (a long-standing default_config contract). + copied = copy.deepcopy(config_data) + return copied if copied else load_config(config_file, **load_kwargs) + +def _ensure_lithops_section(config_data): if 'lithops' not in config_data or not config_data['lithops']: config_data['lithops'] = {} + return config_data['lithops'] - cl = config_data['lithops'] - if 'log_level' not in cl: - cl['log_level'] = c.LOGGER_LEVEL - if 'log_format' not in cl: - cl['log_format'] = c.LOGGER_FORMAT - if 'log_stream' not in cl: - cl['log_stream'] = c.LOGGER_STREAM - if 'log_filename' not in cl: - cl['log_filename'] = None +def _section_with_user_agent(config, backend): + # Falsy sections ({}, None, missing) yield a *new* dict so the original + # config is not mutated. A populated section is updated in place. + section = config[backend] if backend in config and config[backend] else {} + section['user_agent'] = _USER_AGENT + return section - return cl['log_level'], cl['log_format'], cl['log_stream'], cl['log_filename'] +def get_log_info(config_file=None, config_data=None): + """Returns the logging settings of a configuration, filling the defaults""" + config_data = _copy_or_load_config(config_file, config_data, log=False) + lithops_cfg = _ensure_lithops_section(config_data) + + lithops_cfg.setdefault('log_level', c.LOGGER_LEVEL) + lithops_cfg.setdefault('log_format', c.LOGGER_FORMAT) + lithops_cfg.setdefault('log_stream', c.LOGGER_STREAM) + lithops_cfg.setdefault('log_filename', None) + + return ( + lithops_cfg['log_level'], + lithops_cfg['log_format'], + lithops_cfg['log_stream'], + lithops_cfg['log_filename'], + ) -def default_config(config_file=None, config_data=None, config_overwrite={}, load_storage_config=True): + +def _resolve_mode_and_backend(config_data): """ - First checks .lithops_config - then checks LITHOPS_CONFIG_FILE environment variable - then ~/.lithops/config + Fills in the mode and the backend out of each other. When both are set the + backend wins, and the mode is rewritten to the one it belongs to """ - logger.info(f'Lithops v{__version__} - Python{CURRENT_PY_VERSION}') - - config_data = copy.deepcopy(config_data) or load_config(config_file) - - if 'lithops' not in config_data or not config_data['lithops']: - config_data['lithops'] = {} - - # overwrite values provided by the user - if 'lithops' in config_overwrite: - config_data['lithops'].update(config_overwrite['lithops']) - - backend = config_data['lithops'].get('backend') - mode = config_data['lithops'].get('mode') + lithops_cfg = config_data['lithops'] + backend = lithops_cfg.get('backend') + mode = lithops_cfg.get('mode') if mode and not backend: if mode in config_data and 'backend' in config_data[mode]: - config_data['lithops']['backend'] = config_data[mode]['backend'] + lithops_cfg['backend'] = config_data[mode]['backend'] else: - config_data['lithops']['backend'] = get_default_backend(mode) + lithops_cfg['backend'] = get_default_backend(mode) elif backend: - config_data['lithops']['mode'] = get_mode(backend) + lithops_cfg['mode'] = get_mode(backend) elif not backend and not mode: - mode = config_data['lithops']['mode'] = c.MODE_DEFAULT - config_data['lithops']['backend'] = get_default_backend(mode) + mode = lithops_cfg['mode'] = c.MODE_DEFAULT + lithops_cfg['backend'] = get_default_backend(mode) - backend = config_data['lithops'].get('backend') - mode = config_data['lithops'].get('mode') + return lithops_cfg.get('backend'), lithops_cfg.get('mode') - if backend not in config_data or config_data[backend] is None: - config_data[backend] = {} - - if 'backend' in config_overwrite and config_overwrite['backend']: - config_data[backend].update(config_overwrite['backend']) +def _load_compute_backend_config(config_data, mode, backend): + """Lets the config module of the compute backend fill in its own defaults""" if mode == c.LOCALHOST: logger.debug("Loading compute backend module: localhost") - cb_config = importlib.import_module('lithops.localhost.config') - cb_config.load_config(config_data) - + module_name = 'lithops.localhost.config' elif mode == c.SERVERLESS: logger.debug(f"Loading Serverless backend module: {backend}") - cb_config = importlib.import_module(f'lithops.serverless.backends.{backend}.config') - cb_config.load_config(config_data) - + module_name = f'lithops.serverless.backends.{backend}.config' elif mode == c.STANDALONE: logger.debug(f"Loading Standalone backend module: {backend}") - sb_config = importlib.import_module(f'lithops.standalone.backends.{backend}.config') - sb_config.load_config(config_data) + module_name = f'lithops.standalone.backends.{backend}.config' + else: + return + + importlib.import_module(module_name).load_config(config_data) + + if mode == c.STANDALONE: + # Standalone always runs one call per worker; user chunksize is ignored. config_data['lithops']['chunksize'] = 0 + +def default_config( + config_file=None, + config_data=None, + config_overwrite=None, + load_storage_config=True, +): + """ + Build a complete Lithops configuration. + + Config is loaded from `config_data`, `config_file`, or the default file + locations (see `get_default_config_filename`). Values in `config_overwrite` + replace matching keys. + """ + logger.info(f'Lithops v{__version__} - Python{CURRENT_PY_VERSION}') + + config_overwrite = config_overwrite or {} + config_data = _copy_or_load_config(config_file, config_data) + lithops_cfg = _ensure_lithops_section(config_data) + + if 'lithops' in config_overwrite: + lithops_cfg.update(config_overwrite['lithops']) + + backend, mode = _resolve_mode_and_backend(config_data) + + if backend not in config_data or config_data[backend] is None: + # Missing or None is replaced. An existing {} is kept (unlike the + # lithops section, where an empty dict is also replaced). + config_data[backend] = {} + + if 'backend' in config_overwrite and config_overwrite['backend']: + config_data[backend].update(config_overwrite['backend']) + + _load_compute_backend_config(config_data, mode, backend) + if 'chunksize' not in config_data['lithops']: - config_data['lithops']['chunksize'] = config_data[backend]['worker_processes'] + config_data['lithops']['chunksize'] = ( + config_data[backend]['worker_processes'] + ) if load_storage_config: config_data = default_storage_config(config_data=config_data) - if config_data['lithops']['storage'] == c.LOCALHOST and backend != c.LOCALHOST: - raise Exception(f'Localhost storage backend cannot be used with {backend}') + storage = config_data['lithops']['storage'] + if storage == c.LOCALHOST and backend != c.LOCALHOST: + raise Exception( + f'Localhost storage backend cannot be used with {backend}' + ) - for key in c.LITHOPS_DEFAULT_CONFIG_KEYS: - if key not in config_data['lithops']: - config_data['lithops'][key] = c.LITHOPS_DEFAULT_CONFIG_KEYS[key] + for key, value in c.LITHOPS_DEFAULT_CONFIG_KEYS.items(): + config_data['lithops'].setdefault(key, value) return config_data def default_storage_config(config_file=None, config_data=None, backend=None): - """ Function to load default storage config """ - - config_data = copy.deepcopy(config_data) or load_config(config_file) - - if 'lithops' not in config_data or not config_data['lithops']: - config_data['lithops'] = {} + """ + Builds a Lithops configuration that only holds the storage backend, whose + config module fills in its own defaults + """ + config_data = _copy_or_load_config(config_file, config_data) + lithops_cfg = _ensure_lithops_section(config_data) - if 'storage' not in config_data['lithops']: - config_data['lithops']['storage'] = c.STORAGE_BACKEND_DEFAULT + if 'storage' not in lithops_cfg: + lithops_cfg['storage'] = c.STORAGE_BACKEND_DEFAULT if backend: - config_data['lithops']['storage'] = backend + lithops_cfg['storage'] = backend - sb = config_data['lithops']['storage'] + sb = lithops_cfg['storage'] logger.debug(f"Loading Storage backend module: {sb}") - sb_config = importlib.import_module(f'lithops.storage.backends.{sb}.config') - sb_config.load_config(config_data) + importlib.import_module( + f'lithops.storage.backends.{sb}.config' + ).load_config(config_data) return config_data def extract_storage_config(config): - s_config = {} - s_config['monitoring_interval'] = config['lithops'].get( - 'monitoring_interval', c.LITHOPS_DEFAULT_CONFIG_KEYS['monitoring_interval'] - ) + """Extracts the config that the storage backend of a job needs""" backend = config['lithops']['storage'] - s_config['backend'] = backend - s_config[backend] = config[backend] if backend in config and config[backend] else {} - s_config[backend]['user_agent'] = f'lithops/{__version__}' - - return s_config + return { + 'monitoring_interval': config['lithops'].get( + 'monitoring_interval', + c.LITHOPS_DEFAULT_CONFIG_KEYS['monitoring_interval'], + ), + 'backend': backend, + backend: _section_with_user_agent(config, backend), + } def extract_localhost_config(config): - localhost_config = config[c.LOCALHOST].copy() - - return localhost_config + """Extracts the config that the localhost compute backend needs""" + return config[c.LOCALHOST].copy() def extract_serverless_config(config): - sl_config = {} + """Extracts the config that the serverless compute backend needs""" backend = config['lithops']['backend'] - sl_config['backend'] = backend - sl_config[backend] = config[backend] if backend in config and config[backend] else {} - sl_config[backend]['user_agent'] = f'lithops/{__version__}' - - return sl_config + return { + 'backend': backend, + backend: _section_with_user_agent(config, backend), + } def extract_standalone_config(config): - sa_config = config[c.STANDALONE].copy() + """Extracts the config that the standalone compute backend needs""" backend = config['lithops']['backend'] + sa_config = config[c.STANDALONE].copy() sa_config['backend'] = backend sa_config['storage'] = config['lithops'].get('storage') - sa_config[backend] = config[backend] if backend in config and config[backend] else {} - sa_config[backend]['user_agent'] = f'lithops/{__version__}' - + sa_config[backend] = _section_with_user_agent(config, backend) return sa_config diff --git a/lithops/constants.py b/lithops/constants.py index 515899975..02356113d 100644 --- a/lithops/constants.py +++ b/lithops/constants.py @@ -15,19 +15,23 @@ # import os +import posixpath import tempfile +# --------------------------------------------------------------------------- +# Logging +# --------------------------------------------------------------------------- LOGGER_LEVEL = 'info' LOGGER_STREAM = 'ext://sys.stderr' -LOGGER_FORMAT = "%(asctime)s [%(levelname)s] %(filename)s:%(lineno)s -- %(message)s" +LOGGER_FORMAT = ( + "%(asctime)s [%(levelname)s] %(filename)s:%(lineno)s -- %(message)s" +) LOGGER_FORMAT_SHORT = "[%(levelname)s] %(filename)s:%(lineno)s -- %(message)s" -LOGGER_LEVEL_CHOICES = ["debug", "info", "warning", "error", "critical"] - -CPU_COUNT = os.cpu_count() - -STORAGE_CLI_MSG = '{} client created' -COMPUTE_CLI_MSG = '{} client created' +LOGGER_LEVEL_CHOICES = ("debug", "info", "warning", "error", "critical") +# --------------------------------------------------------------------------- +# Execution modes +# --------------------------------------------------------------------------- LOCALHOST = 'localhost' SERVERLESS = 'serverless' STANDALONE = 'standalone' @@ -38,36 +42,53 @@ STANDALONE_BACKEND_DEFAULT = 'aws_ec2' STORAGE_BACKEND_DEFAULT = 'aws_s3' +CPU_COUNT = os.cpu_count() +WORKER_PROCESSES_DEFAULT = 1 +MAX_AGG_DATA_SIZE = 4 # 4MiB + +STORAGE_CLI_MSG = '{} client created' +COMPUTE_CLI_MSG = '{} client created' + +# --------------------------------------------------------------------------- +# Object-storage key prefixes +# --------------------------------------------------------------------------- JOBS_PREFIX = "lithops.jobs" TEMP_PREFIX = "lithops.jobs/tmp" LOGS_PREFIX = "lithops.logs" RUNTIMES_PREFIX = "lithops.runtimes" -MAX_AGG_DATA_SIZE = 4 # 4MiB - -WORKER_PROCESSES_DEFAULT = 1 - +# --------------------------------------------------------------------------- +# Local filesystem +# --------------------------------------------------------------------------- TEMP_DIR = os.path.realpath(tempfile.gettempdir()) -USER_TEMP_DIR = 'lithops-' + os.getenv("USER", "root") +USER_TEMP_DIR = f"lithops-{os.getenv('USER', 'root')}" LITHOPS_TEMP_DIR = os.path.join(TEMP_DIR, USER_TEMP_DIR) -JOBS_DIR = os.path.join(LITHOPS_TEMP_DIR, 'jobs') -LOGS_DIR = os.path.join(LITHOPS_TEMP_DIR, 'logs') -MODULES_DIR = os.path.join(LITHOPS_TEMP_DIR, 'modules') -CUSTOM_RUNTIME_DIR = os.path.join(LITHOPS_TEMP_DIR, 'custom-runtime') -RN_LOG_FILE = os.path.join(LITHOPS_TEMP_DIR, 'localhost-runner.log') -SV_LOG_FILE = os.path.join(LITHOPS_TEMP_DIR, 'localhost-service.log') -FN_LOG_FILE = os.path.join(LITHOPS_TEMP_DIR, 'functions.log') -CLEANER_DIR = os.path.join(LITHOPS_TEMP_DIR, 'cleaner') +def _in_temp(*parts): + return os.path.join(LITHOPS_TEMP_DIR, *parts) + + +JOBS_DIR = _in_temp('jobs') +LOGS_DIR = _in_temp('logs') +MODULES_DIR = _in_temp('modules') +CUSTOM_RUNTIME_DIR = _in_temp('custom-runtime') + +RN_LOG_FILE = _in_temp('localhost-runner.log') +SV_LOG_FILE = _in_temp('localhost-service.log') +FN_LOG_FILE = _in_temp('functions.log') + +CLEANER_DIR = _in_temp('cleaner') CLEANER_PID_FILE = os.path.join(CLEANER_DIR, 'cleaner.pid') CLEANER_LOG_FILE = os.path.join(CLEANER_DIR, 'cleaner.log') +CLEANER_TMP_SUFFIX = '.tmp' HOME_DIR = os.path.expanduser('~') CONFIG_DIR = os.path.join(HOME_DIR, '.lithops') CACHE_DIR = os.path.join(CONFIG_DIR, 'cache') CONFIG_FILE = os.path.join(CONFIG_DIR, 'config') -CONFIG_FILE_GLOBAL = os.path.join("/etc", "lithops", "config") +# /etc is only ever read where Lithops runs on Linux, so this stays POSIX +CONFIG_FILE_GLOBAL = '/etc/lithops/config' LITHOPS_DEFAULT_CONFIG_KEYS = { 'monitoring': 'storage', @@ -75,16 +96,19 @@ 'execution_timeout': 1800 } +# --------------------------------------------------------------------------- +# Standalone VM (remote paths are POSIX — workers run on Linux) +# --------------------------------------------------------------------------- SA_INSTALL_DIR = '/opt/lithops' -SA_SETUP_LOG_FILE = f'{SA_INSTALL_DIR}/setup.log' -SA_SETUP_DONE_FILE = f'{SA_INSTALL_DIR}/setup-done.flag' -SA_MASTER_LOG_FILE = f'{LITHOPS_TEMP_DIR}/master-service.log' -SA_WORKER_LOG_FILE = f'{LITHOPS_TEMP_DIR}/worker-service.log' +SA_SETUP_LOG_FILE = posixpath.join(SA_INSTALL_DIR, 'setup.log') +SA_SETUP_DONE_FILE = posixpath.join(SA_INSTALL_DIR, 'setup-done.flag') +SA_MASTER_LOG_FILE = _in_temp('master-service.log') +SA_WORKER_LOG_FILE = _in_temp('worker-service.log') SA_MASTER_SERVICE_PORT = 8080 SA_WORKER_SERVICE_PORT = 8081 -SA_CONFIG_FILE = os.path.join(SA_INSTALL_DIR, 'config') -SA_MASTER_DATA_FILE = os.path.join(SA_INSTALL_DIR, 'master.data') -SA_WORKER_DATA_FILE = os.path.join(SA_INSTALL_DIR, 'worker.data') +SA_CONFIG_FILE = posixpath.join(SA_INSTALL_DIR, 'config') +SA_MASTER_DATA_FILE = posixpath.join(SA_INSTALL_DIR, 'master.data') +SA_WORKER_DATA_FILE = posixpath.join(SA_INSTALL_DIR, 'worker.data') SA_DEFAULT_CONFIG_KEYS = { 'runtime': 'python3', @@ -98,7 +122,10 @@ 'extra_python_packages': [], } -SERVERLESS_BACKENDS = [ +# --------------------------------------------------------------------------- +# Known compute backends +# --------------------------------------------------------------------------- +SERVERLESS_BACKENDS = ( 'ibm_cf', 'code_engine', 'knative', @@ -113,13 +140,13 @@ 'aliyun_fc', 'oracle_f', 'k8s', - 'singularity' -] + 'singularity', +) -STANDALONE_BACKENDS = [ +STANDALONE_BACKENDS = ( 'ibm_vpc', 'aws_ec2', 'azure_vms', 'gcp_compute_engine', - 'vm' -] + 'vm', +) diff --git a/lithops/executors.py b/lithops/executors.py index 0382d263f..0ec69016d 100644 --- a/lithops/executors.py +++ b/lithops/executors.py @@ -22,7 +22,7 @@ import pickle import tempfile import subprocess as sp -from typing import Optional, List, Union, Tuple, Dict, Any +from typing import Any, Dict, List, Optional, Tuple, Union from collections.abc import Callable from datetime import datetime @@ -30,42 +30,118 @@ from lithops.future import ResponseFuture from lithops.invokers import create_invoker from lithops.storage import InternalStorage -from lithops.wait import wait, ALL_COMPLETED, THREADPOOL_SIZE, ALWAYS +from lithops.wait import ( + wait, + ALL_COMPLETED, + THREADPOOL_SIZE, + ALWAYS, + _partition_futures, +) from lithops.job import create_map_job, create_reduce_job -from lithops.config import default_config, \ - extract_localhost_config, extract_standalone_config, \ - extract_serverless_config, get_log_info, extract_storage_config -from lithops.constants import LOCALHOST, CLEANER_DIR, \ - SERVERLESS, STANDALONE -from lithops.utils import setup_lithops_logger, \ - is_lithops_worker, create_executor_id, create_futures_list +from lithops.job.job import invalidate_function_cache +from lithops.config import ( + default_config, + extract_localhost_config, + extract_standalone_config, + extract_serverless_config, + get_log_info, + extract_storage_config, +) +from lithops.constants import LOCALHOST, CLEANER_DIR, CLEANER_TMP_SUFFIX, SERVERLESS, STANDALONE +from lithops.utils import ( + setup_lithops_logger, + is_lithops_worker, + create_executor_id, + create_futures_list, + FuturesList, + _as_future_list as wrap_as_future_list, + log_prefix, +) from lithops.localhost import LocalhostHandlerV1, LocalhostHandlerV2 from lithops.standalone import StandaloneHandler from lithops.serverless import ServerlessHandler from lithops.storage.utils import create_job_key, CloudObject from lithops.monitor import JobMonitor -from lithops.utils import FuturesList logger = logging.getLogger(__name__) -CLEANER_PROCESS = None + + +def _dump_cleaner_data(data: Dict[str, Any]) -> None: + """ + Drops a request in the shared cleaner directory, which the cleaner + process picks up and deletes once it has honoured it. Every Lithops + process on this machine writes here, and the cleaner reads the directory + while they do, so the request is written under a staging name the + cleaner ignores and then renamed into place: the rename is atomic, and + no cleaner can ever read a half written pickle + """ + os.makedirs(CLEANER_DIR, exist_ok=True) + with tempfile.NamedTemporaryFile( + dir=CLEANER_DIR, suffix=CLEANER_TMP_SUFFIX, delete=False + ) as temp: + pickle.dump(data, temp) + os.replace(temp.name, temp.name[:-len(CLEANER_TMP_SUFFIX)]) + + +def _omit_none(mapping: Dict[str, Any]) -> Dict[str, Any]: + """ + Keeps the entries the user actually set, so that they do not overwrite + the config with None + """ + return {key: value for key, value in mapping.items() if value is not None} + + +def _missing_plotting_extra(method_name: str) -> ModuleNotFoundError: + """ + Error for a method that needs the optional plotting dependencies + """ + return ModuleNotFoundError( + f"Please install 'pip3 install lithops[plotting]' for " + f"making use of the {method_name}() method" + ) + + +def _group_futures_by_job( + futures: List[Any] +) -> List[Tuple[str, str, List[Any], List[Any]]]: + """ + Splits the futures into consecutive runs of the same job, keeping the + execution time and the memory of each of its activations. Every job runs + a single function + """ + groups = [] + for future in futures: + if not groups or groups[-1][0] != future.job_id: + groups.append((future.job_id, future.function_name, [], [])) + _, _, runtimes, memory = groups[-1] + runtimes.append(future.stats['worker_exec_time']) + memory.append(future.runtime_memory) + return groups class FunctionExecutor: """ - Executor abstract class that contains the common logic for the Localhost, Serverless and Standalone executors + Base executor that contains the common logic for the Localhost, Serverless + and Standalone executors. :param mode: Execution mode. One of: localhost, serverless or standalone :param config: Settings passed in here will override those in lithops_config :param config_file: Path to the lithops config file :param backend: Compute backend to run the functions :param storage: Storage backend to store Lithops data - :param monitoring: Monitoring system implementation. One of: storage, rabbitmq - :param log_level: Log level printing (INFO, DEBUG, ...). Set it to None to hide all logs. - If this is param is set, all logging params in config are disabled - :param kwargs: Any parameter that can be set in the compute backend section of the config file, can be set here + :param monitoring: Monitoring system implementation. + One of: storage, rabbitmq + :param log_level: Log level printing (INFO, DEBUG, ...). + Set it to None to hide all logs. + If this is param is set, all logging params in config + are disabled + :param kwargs: Any parameter that can be set in the compute + backend section of the config file, can be set here """ + _cleaner_process = None + def __init__( self, mode: Optional[str] = None, @@ -74,8 +150,8 @@ def __init__( backend: Optional[str] = None, storage: Optional[str] = None, monitoring: Optional[str] = None, - log_level: Optional[str] = False, - **kwargs: Optional[Dict[str, Any]] + log_level: Union[str, bool, None] = False, + **kwargs: Any ): self.is_lithops_worker = is_lithops_worker() self.executor_id = create_executor_id() @@ -83,32 +159,26 @@ def __init__( self.cleaned_jobs = set() self.total_jobs = 0 self.last_call = None + self.log_path = None - # setup lithops logging - if not self.is_lithops_worker: - # if is lithops worker, logging has been set up in entry_point.py - if log_level: - setup_lithops_logger(log_level) - elif log_level is False and logger.getEffectiveLevel() == logging.WARNING: - # Set default logging from config - setup_lithops_logger(*get_log_info(config_file=config_file, config_data=config)) - - # overwrite user-provided parameters - config_ow = {'lithops': {}, 'backend': {}} - for key, value in kwargs.items(): - if value is not None: - config_ow['backend'][key] = value - args = {'mode': mode, 'backend': backend, 'storage': storage, 'monitoring': monitoring} - for key, value in args.items(): - if value is not None: - config_ow['lithops'][key] = value - - # Load configuration - self.config = default_config(config_file=config_file, config_data=config, config_overwrite=config_ow) + self._setup_logging(log_level, config_file, config) + + self.config = default_config( + config_file=config_file, + config_data=config, + config_overwrite=self._build_config_overwrite( + mode, backend, storage, monitoring, kwargs + ) + ) self.data_cleaner = self.config['lithops'].get('data_cleaner', True) if self.data_cleaner and not self.is_lithops_worker: - atexit.register(self.clean, clean_cloudobjects=False, clean_fn=True, on_exit=True) + atexit.register( + self.clean, + clean_cloudobjects=False, + clean_fn=True, + on_exit=True, + ) storage_config = extract_storage_config(self.config) self.internal_storage = InternalStorage(storage_config) @@ -116,30 +186,16 @@ def __init__( self.backend = self.config['lithops']['backend'] self.mode = self.config['lithops']['mode'] + self.compute_handler = self._create_compute_handler() + self.config['lithops']['backend_type'] = ( + self.compute_handler.get_backend_type() + ) - if self.mode == LOCALHOST: - localhost_config = extract_localhost_config(self.config) - if localhost_config.get('version', 2) == 1: - self.compute_handler = LocalhostHandlerV1(localhost_config) - else: - self.compute_handler = LocalhostHandlerV2(localhost_config) - elif self.mode == SERVERLESS: - serverless_config = extract_serverless_config(self.config) - self.compute_handler = ServerlessHandler(serverless_config, self.internal_storage) - elif self.mode == STANDALONE: - standalone_config = extract_standalone_config(self.config) - self.compute_handler = StandaloneHandler(standalone_config) - - self.config['lithops']['backend_type'] = self.compute_handler.get_backend_type() - - # Create the monitoring system self.job_monitor = JobMonitor( executor_id=self.executor_id, internal_storage=self.internal_storage, config=self.config ) - - # Create the invoker self.invoker = create_invoker( config=self.config, executor_id=self.executor_id, @@ -148,25 +204,225 @@ def __init__( job_monitor=self.job_monitor ) - logger.debug(f'Function executor for {self.backend} created with ID: {self.executor_id}') - - self.log_path = None + logger.debug( + f'Function executor for {self.backend} created with ID: {self.executor_id}' + ) def __enter__(self): - """ Context manager method """ + """Context manager method.""" return self def __exit__(self, exc_type, exc_value, traceback): - """ Context manager method """ + """Context manager method.""" self.job_monitor.stop() self.invoker.stop() self.compute_handler.clear() + @staticmethod + def _build_config_overwrite(mode, backend, storage, monitoring, kwargs): + """ + Turns the arguments of the constructor into the config overwrite that + takes precedence over the config file + """ + return { + 'lithops': _omit_none({ + 'mode': mode, + 'backend': backend, + 'storage': storage, + 'monitoring': monitoring, + }), + 'backend': _omit_none(kwargs), + } + + def _setup_logging(self, log_level, config_file, config): + """ + Sets up the Lithops logger, unless the user already configured a + logger of their own or asked for no logs at all + """ + if self.is_lithops_worker: + # Logging has already been set up in entry_point.py + return + if log_level: + setup_lithops_logger(log_level) + elif ( + log_level is False + and logger.getEffectiveLevel() == logging.WARNING + ): + setup_lithops_logger( + *get_log_info( + config_file=config_file, config_data=config + ) + ) + + def _create_compute_handler(self): + """ + Builds the handler of the backend this executor runs on + """ + if self.mode == LOCALHOST: + localhost_config = extract_localhost_config(self.config) + if localhost_config.get('version', 2) == 1: + return LocalhostHandlerV1(localhost_config) + return LocalhostHandlerV2(localhost_config) + if self.mode == SERVERLESS: + return ServerlessHandler( + extract_serverless_config(self.config), + self.internal_storage + ) + if self.mode == STANDALONE: + return StandaloneHandler(extract_standalone_config(self.config)) + return None + def _create_job_id(self, call_type): + """ + Numbers a new job of this executor, prefixed by the call that + submitted it: A for call_async, M for map and R for reduce + """ job_id = str(self.total_jobs).zfill(3) self.total_jobs += 1 return f'{call_type}{job_id}' + @staticmethod + def _as_future_list(futures): + """Keep list subclasses (including FuturesList) unchanged; wrap a single future.""" + return wrap_as_future_list(futures) + + @staticmethod + def _disable_iterdata_output(iterdata): + """ + Marks the futures used as input as consumed, so that get_result() + returns the output of this job only + """ + if isinstance(iterdata, FuturesList): + for fut in iterdata: + fut._produce_output = False + + def _invoke(self, job): + """ + Invokes a job and tracks its futures in this executor + """ + futures = self.invoker.run_job(job) + self.futures.extend(futures) + return futures + + def _run_map_job( + self, + job_id, + map_function, + iterdata, + runtime_memory=None, + **job_kwargs, + ): + """ + Builds a map job and invokes it, returning the job and its futures + """ + runtime_meta = self.invoker.select_runtime(job_id, runtime_memory) + job = create_map_job( + config=self.config, + internal_storage=self.internal_storage, + executor_id=self.executor_id, + job_id=job_id, + map_function=map_function, + iterdata=iterdata, + runtime_meta=runtime_meta, + runtime_memory=runtime_memory, + **job_kwargs + ) + return job, self._invoke(job) + + def _run_reduce_job( + self, + reduce_job_id, + reduce_function, + map_job, + map_futures, + runtime_memory=None, + **job_kwargs, + ): + """ + Builds a reduce job over the futures of a map job and invokes it + """ + runtime_meta = self.invoker.select_runtime( + reduce_job_id, runtime_memory + ) + job = create_reduce_job( + config=self.config, + internal_storage=self.internal_storage, + executor_id=self.executor_id, + reduce_job_id=reduce_job_id, + reduce_function=reduce_function, + map_job=map_job, + map_futures=map_futures, + runtime_meta=runtime_meta, + runtime_memory=runtime_memory, + **job_kwargs + ) + return self._invoke(job) + + def _submit_map( + self, + map_function, + iterdata, + runtime_memory=None, + extra_env=None, + include_modules=None, + exclude_modules=None, + timeout=None, + chunksize=None, + extra_args=None, + obj_chunk_size=None, + obj_chunk_number=None, + obj_newline='\n', + job_prefix='M' + ): + """ + Common path of call_async(), map() and the map stage of map_reduce() + """ + job_id = self._create_job_id(job_prefix) + job, futures = self._run_map_job( + job_id=job_id, + map_function=map_function, + iterdata=iterdata, + runtime_memory=runtime_memory, + extra_env=extra_env, + include_modules=include_modules, + exclude_modules=exclude_modules, + execution_timeout=timeout, + chunksize=chunksize, + extra_args=extra_args, + obj_chunk_size=obj_chunk_size, + obj_chunk_number=obj_chunk_number, + obj_newline=obj_newline + ) + self._disable_iterdata_output(iterdata) + return job_id, job, futures + + def _cleanup_jobs(self, futures, exception=None, force=False): + """ + Releases the backend resources of the given jobs and deletes their + temporary data. Not every backend takes the exception that ended them + """ + present_jobs = {f.job_key for f in futures} + if exception is None: + self.compute_handler.clear(present_jobs) + else: + self.compute_handler.clear(present_jobs, exception=exception) + self.clean(clean_cloudobjects=False, force=force) + + def _stop_monitor_if_idle(self, extra_fs=None): + """ + Stops the job monitor once there is no future left to watch, counting + the ones that do not belong to this executor + """ + tracked = list(self.futures) + if extra_fs: + seen = {id(fut) for fut in tracked} + tracked.extend(fut for fut in extra_fs if id(fut) not in seen) + if tracked and all( + getattr(fut, 'ready', False) or fut.success or fut.done + for fut in tracked + ): + self.job_monitor.stop() + def call_async( self, func: Callable, @@ -181,44 +437,45 @@ def call_async( For running one function execution asynchronously. :param func: The function to map over the data. - :param data: Input data. Arguments can be passed as a list or tuple, or as a dictionary for keyword arguments. - :param extra_env: Additional env variables for function environment. + :param data: Input data. Arguments can be passed as a + list or tuple, or as a dictionary for keyword + arguments. + :param extra_env: Additional env variables for function + environment. :param runtime_memory: Memory to use to run the function. - :param timeout: Time that the function has to complete its execution before raising a timeout. - :param include_modules: Explicitly pickle these dependencies. - :param exclude_modules: Explicitly keep these modules from pickled dependencies. + :param timeout: Time that the function has to complete + its execution before raising a timeout. + :param include_modules: Explicitly pickle these + dependencies. + :param exclude_modules: Explicitly keep these modules + from pickled dependencies. :return: Response future. """ - job_id = self._create_job_id('A') self.last_call = 'call_async' - - runtime_meta = self.invoker.select_runtime(job_id, runtime_memory) - - job = create_map_job(config=self.config, - internal_storage=self.internal_storage, - executor_id=self.executor_id, - job_id=job_id, - map_function=func, - iterdata=[data], - runtime_meta=runtime_meta, - runtime_memory=runtime_memory, - extra_env=extra_env, - include_modules=include_modules, - exclude_modules=exclude_modules, - execution_timeout=timeout) - - futures = self.invoker.run_job(job) - self.futures.extend(futures) + _, _, futures = self._submit_map( + func, + [data], + runtime_memory=runtime_memory, + extra_env=extra_env, + include_modules=include_modules, + exclude_modules=exclude_modules, + timeout=timeout, + job_prefix='A' + ) return futures[0] def map( self, map_function: Callable, - map_iterdata: List[Union[List[Any], Tuple[Any, ...], Dict[str, Any]]], + map_iterdata: List[Union[ + List[Any], Tuple[Any, ...], Dict[str, Any] + ]], chunksize: Optional[int] = None, - extra_args: Optional[Union[List[Any], Tuple[Any, ...], Dict[str, Any]]] = None, + extra_args: Optional[Union[ + List[Any], Tuple[Any, ...], Dict[str, Any] + ]] = None, extra_env: Optional[Dict[str, str]] = None, runtime_memory: Optional[int] = None, obj_chunk_size: Optional[int] = None, @@ -229,70 +486,78 @@ def map( exclude_modules: Optional[List[str]] = [] ) -> FuturesList: """ - Spawn multiple function activations based on the items of an input list. + Spawn multiple function activations based on the items + of an input list. :param map_function: The function to map over the data - :param map_iterdata: An iterable of input data (e.g python list). - :param chunksize: Split map_iteradata in chunks of this size. Lithops spawns 1 worker per resulting chunk - :param extra_args: Additional arguments to pass to each map_function activation - :param extra_env: Additional environment variables for function environment - :param runtime_memory: Memory (in MB) to use to run the functions - :param obj_chunk_size: Used for data processing. Chunk size to split each object in bytes. - Must be >= 1MiB. 'None' for processing the whole file in one function activation - :param obj_chunk_number: Used for data processing. Number of chunks to split each object. - 'None' for processing the whole file in one function activation. chunk_n has prevalence over chunk_size if both parameters are set - :param obj_newline: new line character for keeping line integrity of partitions. - 'None' for disabling line integrity logic and get partitions of the exact same size in the functions + :param map_iterdata: An iterable of input data + (e.g python list). + :param chunksize: Split map_iterdata in chunks of this + size. Lithops spawns 1 worker per resulting chunk + :param extra_args: Additional arguments to pass to each + map_function activation + :param extra_env: Additional environment variables for + function environment + :param runtime_memory: Memory (in MB) to use to run + the functions + :param obj_chunk_size: Used for data processing. Chunk + size to split each object in bytes. Must be >= 1MiB. + 'None' for processing the whole file in one + function activation + :param obj_chunk_number: Used for data processing. Number + of chunks to split each object. 'None' for processing + the whole file in one function activation. chunk_n + has prevalence over chunk_size if both parameters + are set + :param obj_newline: new line character for keeping line + integrity of partitions. 'None' for disabling line + integrity logic and get partitions of the exact same + size in the functions :param timeout: Max time per function activation (seconds) - :param include_modules: Explicitly pickle these dependencies. All required dependencies are pickled if default empty list. - No one dependency is pickled if it is explicitly set to None - :param exclude_modules: Explicitly keep these modules from pickled dependencies. It is not taken into account if you set include_modules. - - :return: A list with size `len(map_iterdata)` of futures for each job (Futures are also internally stored by Lithops). + :param include_modules: Explicitly pickle these + dependencies. All required dependencies are pickled + if default empty list. No one dependency is pickled + if it is explicitly set to None + :param exclude_modules: Explicitly keep these modules + from pickled dependencies. It is not taken into + account if you set include_modules. + + :return: A list with size `len(map_iterdata)` of futures + for each job (Futures are also internally stored + by Lithops). """ - - job_id = self._create_job_id('M') self.last_call = 'map' - - runtime_meta = self.invoker.select_runtime(job_id, runtime_memory) - - job = create_map_job( - config=self.config, - internal_storage=self.internal_storage, - executor_id=self.executor_id, - job_id=job_id, - map_function=map_function, - iterdata=map_iterdata, - chunksize=chunksize, - runtime_meta=runtime_meta, + _, _, futures = self._submit_map( + map_function, + map_iterdata, runtime_memory=runtime_memory, extra_env=extra_env, include_modules=include_modules, exclude_modules=exclude_modules, - execution_timeout=timeout, + timeout=timeout, + chunksize=chunksize, extra_args=extra_args, obj_chunk_size=obj_chunk_size, obj_chunk_number=obj_chunk_number, obj_newline=obj_newline ) - futures = self.invoker.run_job(job) - self.futures.extend(futures) - - if isinstance(map_iterdata, FuturesList): - for fut in map_iterdata: - fut._produce_output = False - return create_futures_list(futures, self) def map_reduce( self, map_function: Callable, - map_iterdata: List[Union[List[Any], Tuple[Any, ...], Dict[str, Any]]], + map_iterdata: List[Union[ + List[Any], Tuple[Any, ...], Dict[str, Any] + ]], reduce_function: Callable, chunksize: Optional[int] = None, - extra_args: Optional[Union[List[Any], Tuple[Any, ...], Dict[str, Any]]] = None, - extra_args_reduce: Optional[Union[List[Any], Tuple[Any, ...], Dict[str, Any]]] = None, + extra_args: Optional[Union[ + List[Any], Tuple[Any, ...], Dict[str, Any] + ]] = None, + extra_args_reduce: Optional[Union[ + List[Any], Tuple[Any, ...], Dict[str, Any] + ]] = None, extra_env: Optional[Dict[str, str]] = None, map_runtime_memory: Optional[int] = None, reduce_runtime_memory: Optional[int] = None, @@ -306,79 +571,78 @@ def map_reduce( exclude_modules: Optional[List[str]] = [] ) -> FuturesList: """ - Map the map_function over the data and apply the reduce_function across all futures. + Map the map_function over the data and apply the + reduce_function across all futures. :param map_function: The function to map over the data :param map_iterdata: An iterable of input data - :param reduce_function: The function to reduce over the futures - :param chunksize: Split map_iteradata in chunks of this size. Lithops spawns 1 worker per resulting chunk. Default 1 - :param extra_args: Additional arguments to pass to function activation. Default None - :param extra_args_reduce: Additional arguments to pass to the reduce function activation. Default None - :param extra_env: Additional environment variables for action environment. Default None - :param map_runtime_memory: Memory to use to run the map function. Default None (loaded from config) - :param reduce_runtime_memory: Memory to use to run the reduce function. Default None (loaded from config) - :param timeout: Time that the functions have to complete their execution before raising a timeout - :param obj_chunk_size: the size of the data chunks to split each object. 'None' for processing the whole file in one function activation - :param obj_chunk_number: Number of chunks to split each object. 'None' for processing the whole file in one function activation - :param obj_newline: New line character for keeping line integrity of partitions. - 'None' for disabling line integrity logic and get partitions of the exact same size in the functions - :param obj_reduce_by_key: Set one reducer per object after running the partitioner. By default there is one reducer for all the objects - :param spawn_reducer: Percentage of done map functions before spawning the reduce function - :param include_modules: Explicitly pickle these dependencies. - :param exclude_modules: Explicitly keep these modules from pickled dependencies. + :param reduce_function: The function to reduce over + the futures + :param chunksize: Split map_iterdata in chunks of this + size. Lithops spawns 1 worker per resulting chunk. + Default 1 + :param extra_args: Additional arguments to pass to + function activation. Default None + :param extra_args_reduce: Additional arguments to pass + to the reduce function activation. Default None + :param extra_env: Additional environment variables for + action environment. Default None + :param map_runtime_memory: Memory to use to run the map + function. Default None (loaded from config) + :param reduce_runtime_memory: Memory to use to run the + reduce function. Default None (loaded from config) + :param timeout: Time that the functions have to complete + their execution before raising a timeout + :param obj_chunk_size: the size of the data chunks to + split each object. 'None' for processing the whole + file in one function activation + :param obj_chunk_number: Number of chunks to split each + object. 'None' for processing the whole file in one + function activation + :param obj_newline: New line character for keeping line + integrity of partitions. 'None' for disabling line + integrity logic and get partitions of the exact same + size in the functions + :param obj_reduce_by_key: Set one reducer per object + after running the partitioner. By default there is + one reducer for all the objects + :param spawn_reducer: Percentage of done map functions + before spawning the reduce function + :param include_modules: Explicitly pickle these + dependencies. + :param exclude_modules: Explicitly keep these modules + from pickled dependencies. :return: A list with size `len(map_iterdata)` of futures. """ self.last_call = 'map_reduce' - map_job_id = self._create_job_id('M') - - runtime_meta = self.invoker.select_runtime(map_job_id, map_runtime_memory) - - map_job = create_map_job( - config=self.config, - internal_storage=self.internal_storage, - executor_id=self.executor_id, - job_id=map_job_id, - map_function=map_function, - iterdata=map_iterdata, - chunksize=chunksize, - runtime_meta=runtime_meta, + map_job_id, map_job, map_futures = self._submit_map( + map_function, + map_iterdata, runtime_memory=map_runtime_memory, - extra_args=extra_args, extra_env=extra_env, - obj_chunk_size=obj_chunk_size, - obj_chunk_number=obj_chunk_number, - obj_newline=obj_newline, include_modules=include_modules, exclude_modules=exclude_modules, - execution_timeout=timeout + timeout=timeout, + chunksize=chunksize, + extra_args=extra_args, + obj_chunk_size=obj_chunk_size, + obj_chunk_number=obj_chunk_number, + obj_newline=obj_newline ) - map_futures = self.invoker.run_job(map_job) - self.futures.extend(map_futures) - - if isinstance(map_iterdata, FuturesList): - for fut in map_iterdata: - fut._produce_output = False - if spawn_reducer != ALWAYS: self.wait(map_futures, return_when=spawn_reducer) - logger.debug(f'ExecutorID {self.executor_id} | JobID {map_job_id} - ' - f'{spawn_reducer}% of map activations done. Spawning reduce stage') - - reduce_job_id = map_job_id.replace('M', 'R') - - runtime_meta = self.invoker.select_runtime(reduce_job_id, reduce_runtime_memory) + logger.debug( + f'{log_prefix(self.executor_id, map_job_id)} - {spawn_reducer}% of map ' + f'activations done. Spawning reduce stage' + ) - reduce_job = create_reduce_job( - config=self.config, - internal_storage=self.internal_storage, - executor_id=self.executor_id, - reduce_job_id=reduce_job_id, + reduce_futures = self._run_reduce_job( + reduce_job_id=map_job_id.replace('M', 'R'), reduce_function=reduce_function, map_job=map_job, map_futures=map_futures, - runtime_meta=runtime_meta, runtime_memory=reduce_runtime_memory, extra_args=extra_args_reduce, obj_reduce_by_key=obj_reduce_by_key, @@ -387,16 +651,16 @@ def map_reduce( exclude_modules=exclude_modules ) - reduce_futures = self.invoker.run_job(reduce_job) - self.futures.extend(reduce_futures) - - [f._set_mapreduce() for f in map_futures] + for future in map_futures: + future._set_mapreduce() return create_futures_list(map_futures + reduce_futures, self) def wait( self, - fs: Optional[Union[ResponseFuture, FuturesList, List[ResponseFuture]]] = None, + fs: Optional[Union[ + ResponseFuture, FuturesList, List[ResponseFuture] + ]] = None, throw_except: Optional[bool] = True, return_when: Optional[Any] = ALL_COMPLETED, download_results: Optional[bool] = False, @@ -406,70 +670,76 @@ def wait( show_progressbar: Optional[bool] = True ) -> Tuple[FuturesList, FuturesList]: """ - Wait for the Future instances (possibly created by different Executor instances) - given by fs to complete. Returns a named 2-tuple of sets. The first set, named done, - contains the futures that completed (finished or cancelled futures) before the wait - completed. The second set, named not_done, contains the futures that did not complete - (pending or running futures). timeout can be used to control the maximum number of - seconds to wait before returning. + Wait for the Future instances (possibly created by + different Executor instances) given by fs to complete. + Returns a named 2-tuple of sets. The first set, named + done, contains the futures that completed (finished or + cancelled futures) before the wait completed. The second + set, named not_done, contains the futures that did not + complete (pending or running futures). timeout can be + used to control the maximum number of seconds to wait + before returning. :param fs: Futures list. Default None - :param throw_except: Re-raise exception if call raised. Default True + :param throw_except: Re-raise exception if call raised. + Default True :param return_when: Percentage of done futures - :param download_results: Download results. Default false (Only get statuses) + :param download_results: Download results. Default false + (Only get statuses) :param timeout: Timeout of waiting for results - :param threadpool_size: Number of threads to use. Default 64 - :param wait_dur_sec: Time interval between each check. Default 1 second - :param show_progressbar: whether or not to show the progress bar. - - :return: `(fs_done, fs_notdone)` where `fs_done` is a list of futures that have - completed and `fs_notdone` is a list of futures that have not completed. + :param threadpool_size: Number of threads to use. + Default 64 + :param wait_dur_sec: Time interval between each check. + Default 1 second + :param show_progressbar: whether or not to show the + progress bar. + + :return: `(fs_done, fs_notdone)` where `fs_done` is a + list of futures that have completed and `fs_notdone` + is a list of futures that have not completed. """ - futures = fs or self.futures - - if type(futures) not in [list, FuturesList]: - futures = [futures] + futures = self._as_future_list(fs or self.futures) try: - wait(fs=futures, - internal_storage=self.internal_storage, - job_monitor=self.job_monitor, - download_results=download_results, - throw_except=throw_except, - return_when=return_when, - timeout=timeout, - threadpool_size=threadpool_size, - wait_dur_sec=wait_dur_sec, - show_progressbar=show_progressbar, - futures_from_executor_wait=False if fs else True) + wait( + fs=futures, + internal_storage=self.internal_storage, + job_monitor=self.job_monitor, + download_results=download_results, + throw_except=throw_except, + return_when=return_when, + timeout=timeout, + threadpool_size=threadpool_size, + wait_dur_sec=wait_dur_sec, + show_progressbar=show_progressbar, + futures_from_executor_wait=not fs, + ) if self.data_cleaner and return_when == ALL_COMPLETED: - present_jobs = {f.job_key for f in futures} - self.compute_handler.clear(present_jobs) - self.clean(clean_cloudobjects=False) + self._cleanup_jobs(futures) + self._stop_monitor_if_idle(futures) except (KeyboardInterrupt, Exception) as e: self.invoker.stop() self.job_monitor.remove(futures) - [f._set_exception() for f in futures] + for future in futures: + future._set_exception() if self.data_cleaner: - present_jobs = {f.job_key for f in futures} - self.compute_handler.clear(present_jobs, exception=e) - self.clean(clean_cloudobjects=False, force=True) - raise e - - if download_results: - fs_done = [f for f in futures if f.done] - fs_notdone = [f for f in futures if not f.done] - else: - fs_done = [f for f in futures if f.success or f.done] - fs_notdone = [f for f in futures if not f.success and not f.done] - - return create_futures_list(fs_done, self), create_futures_list(fs_notdone, self) + self._cleanup_jobs(futures, exception=e, force=True) + self._stop_monitor_if_idle(futures) + raise + + fs_done, fs_notdone = _partition_futures(futures, download_results) + return ( + create_futures_list(fs_done, self), + create_futures_list(fs_notdone, self), + ) def get_result( self, - fs: Optional[Union[ResponseFuture, FuturesList, List[ResponseFuture]]] = None, + fs: Optional[Union[ + ResponseFuture, FuturesList, List[ResponseFuture] + ]] = None, throw_except: Optional[bool] = True, timeout: Optional[int] = None, threadpool_size: Optional[int] = THREADPOOL_SIZE, @@ -482,18 +752,20 @@ def get_result( :param fs: Futures list. Default None :param throw_except: Reraise exception if call raised. Default True. :param timeout: Timeout for waiting for results. - :param threadpool_size: Number of threads to use. Default 128 + :param threadpool_size: Number of threads to use. Default 64 :param wait_dur_sec: Time interval between each check. Default 1 second :param show_progressbar: whether or not to show the progress bar. :return: The result of the future/s """ - pending_to_read = len(fs) if fs else len( - [f for f in self.futures if not f._read and not f.futures]) + pending_to_read = ( + len(fs) if fs + else sum(1 for f in self.futures if not f._read and not f.futures) + ) logger.info( - (f'ExecutorID {self.executor_id} - Getting results from ' - f'{pending_to_read} function activations') + f'{log_prefix(self.executor_id)} - Getting results from ' + f'{pending_to_read} function activations' ) fs_done, _ = self.wait( @@ -507,16 +779,21 @@ def get_result( ) result = [] - for f in [f for f in fs_done if not f.futures and f._produce_output]: - if fs: # Process futures provided by the user - result.append(f.result(throw_except=throw_except, - internal_storage=self.internal_storage)) - elif not fs and not f._read: # Process internally stored futures - result.append(f.result(throw_except=throw_except, - internal_storage=self.internal_storage)) - f._read = True + for future in fs_done: + if future.futures or not future._produce_output: + continue + if not fs and future._read: + continue + result.append(future.result( + throw_except=throw_except, + internal_storage=self.internal_storage + )) + if not fs: + future._read = True - logger.debug(f'ExecutorID {self.executor_id} - Finished getting results') + logger.debug( + f'{log_prefix(self.executor_id)} - Finished getting results' + ) if len(result) == 1 and self.last_call != 'map': return result[0] @@ -525,40 +802,63 @@ def get_result( def plot( self, - fs: Optional[Union[ResponseFuture, List[ResponseFuture], FuturesList]] = None, + fs: Optional[Union[ + ResponseFuture, List[ResponseFuture], FuturesList + ]] = None, dst: Optional[str] = None, figsize: Optional[tuple] = (10, 6) ): """ - Creates timeline and histogram of the current execution in dst_dir. + Creates timeline and histogram of the current execution in dst. :param fs: list of futures. :param dst: destination path to save .png plots. + :param figsize: size of the plots, in inches. """ - ftrs = self.futures if not fs else fs - + ftrs = fs or self.futures if isinstance(ftrs, ResponseFuture): ftrs = [ftrs] - ftrs_to_plot = [f for f in ftrs if (f.success or f.done) and not f.error] + ftrs_to_plot = [ + f for f in ftrs + if (f.success or f.done) and not f.error + ] if not ftrs_to_plot: - logger.debug(f'ExecutorID {self.executor_id} - No futures ready to plot') + logger.debug( + f'{log_prefix(self.executor_id)} - No futures ready to plot' + ) return try: logging.getLogger('matplotlib').setLevel(logging.WARNING) from lithops.plots import create_timeline, create_histogram except ImportError: - raise ModuleNotFoundError( - "Please install 'pip3 install lithops[plotting]' for " - "making use of the plot() method") + raise _missing_plotting_extra('plot') - logger.info(f'ExecutorID {self.executor_id} - Creating execution plots') + logger.info(f'{log_prefix(self.executor_id)} - Creating execution plots') create_timeline(ftrs_to_plot, dst, figsize) create_histogram(ftrs_to_plot, dst, figsize) + @staticmethod + def _spawn_cleaner_process(): + """ + Starts the process that honours the pending cleaner requests. One + cleaner picks up every request, so a running one is left alone + """ + cleaner = FunctionExecutor._cleaner_process + if cleaner and cleaner.poll() is None: + return + + FunctionExecutor._cleaner_process = sp.Popen( + [sys.executable, '-m', 'lithops.scripts.cleaner'], + start_new_session=True, + env=os.environ.copy(), + stdout=sp.DEVNULL, + stderr=sp.DEVNULL + ) + def clean( self, fs: Optional[Union[ResponseFuture, List[ResponseFuture]]] = None, @@ -569,155 +869,132 @@ def clean( on_exit: Optional[bool] = False ): """ - Deletes all the temp files from storage. These files include the function, - the data serialization and the function invocation results. It can also clean + Deletes all the temp files from storage. These files + include the function, the data serialization and the + function invocation results. It can also clean cloudobjects. :param fs: List of futures to clean :param cs: List of cloudobjects to clean - :param clean_cloudobjects: Delete all cloudobjects created with this executor + :param clean_cloudobjects: Delete all cloudobjects + created with this executor :param clean_fn: Delete cached functions in this executor - :param force: Clean all future objects even if they have not benn completed - :parma on_exit: do not print logs on exit + :param force: Clean all future objects even if they have + not been completed + :param on_exit: do not print logs on exit """ - global CLEANER_PROCESS - - def save_data_to_clean(data): - with tempfile.NamedTemporaryFile(dir=CLEANER_DIR, delete=False) as temp: - pickle.dump(data, temp) - - try: - self.internal_storage - except AttributeError: + if not hasattr(self, 'internal_storage'): return + storage_config = self.internal_storage.get_storage_config() + if cs: - data = { + _dump_cleaner_data({ 'cos_to_clean': list(cs), - 'storage_config': self.internal_storage.get_storage_config() - } - save_data_to_clean(data) + 'storage_config': storage_config + }) if not fs: return if clean_fn: - data = { + invalidate_function_cache(self.executor_id) + _dump_cleaner_data({ 'fn_to_clean': self.executor_id, - 'storage_config': self.internal_storage.get_storage_config() - } - save_data_to_clean(data) - - futures = fs or self.futures - futures = [futures] if type(futures) is not list else futures - present_jobs = {create_job_key(f.executor_id, f.job_id) for f in futures - if (f.executor_id.count('-') == 1 and f.done) or force} + 'storage_config': storage_config + }) + + futures = self._as_future_list(fs or self.futures) + present_jobs = { + create_job_key(f.executor_id, f.job_id) + for f in futures + if (f.executor_id.count('-') == 1 and f.done) or force + } jobs_to_clean = present_jobs - self.cleaned_jobs if jobs_to_clean: if not on_exit: - logger.info(f'ExecutorID {self.executor_id} - Cleaning temporary data') - data = { + logger.info( + f'{log_prefix(self.executor_id)} - Cleaning temporary data' + ) + _dump_cleaner_data({ 'jobs_to_clean': jobs_to_clean, 'clean_cloudobjects': clean_cloudobjects, - 'storage_config': self.internal_storage.get_storage_config() - } - save_data_to_clean(data) + 'storage_config': storage_config + }) self.cleaned_jobs.update(jobs_to_clean) - spawn_cleaner = not (CLEANER_PROCESS and CLEANER_PROCESS.poll() is None) - if (jobs_to_clean or cs) and spawn_cleaner: - cmd = [sys.executable, '-m', 'lithops.scripts.cleaner'] - env = os.environ.copy() - CLEANER_PROCESS = sp.Popen( - cmd, - start_new_session=True, - env=env, - stdout=sp.DEVNULL, - stderr=sp.DEVNULL - ) + if jobs_to_clean or cs: + self._spawn_cleaner_process() def job_summary(self, cloud_objects_n: Optional[int] = 0): """ - Logs information of a job executed by the calling function executor. - currently supports: code_engine, ibm_vpc and ibm_cf. + Logs information of a job executed by the calling + function executor. currently supports: code_engine, + ibm_vpc and ibm_cf. - :param cloud_objects_n: number of cloud object used in COS, declared by user. + :param cloud_objects_n: number of cloud object used in + COS, declared by user. """ try: import pandas as pd import numpy as np except ImportError: - raise ModuleNotFoundError( - "Please install 'pip3 install lithops[plotting]' for " - "making use of the job_summary() method") - - def init(): - headers = ['Job_ID', 'Function', 'Invocations', 'Memory(MB)', 'AvgRuntime', 'Cost', 'CloudObjects'] - pd.DataFrame([], columns=headers).to_csv(self.log_path, index=False) - - def append(content): - """ appends job information to log file.""" - pd.DataFrame(content).to_csv(self.log_path, mode='a', header=False, index=False) - - def append_summary(): - """ add a summary row to the log file""" - df = pd.read_csv(self.log_path) - total_average = sum(df.AvgRuntime * df.Invocations) / df.Invocations.sum() - total_row = pd.DataFrame([['Summary', ' ', df.Invocations.sum(), df['Memory(MB)'].sum(), - round(total_average, 10), df.Cost.sum(), cloud_objects_n]]) - total_row.to_csv(self.log_path, mode='a', header=False, index=False) - - def get_object_num(): - """returns cloud objects used up to this point, using this function executor. """ - df = pd.read_csv(self.log_path) - return float(df.iloc[-1].iloc[-1]) - - # Avoid logging info unless chosen computational backend is supported. - if hasattr(self.compute_handler.backend, 'calc_cost'): - - if self.log_path: # retrieve cloud_objects_n from last log file - cloud_objects_n += get_object_num() - else: - self.log_path = os.path.join(constants.LOGS_DIR, datetime.now().strftime("%Y-%m-%d_%H-%M-%S.csv")) - # override current logfile - init() - - futures = self.futures - if type(futures) is not list: - futures = [futures] - - memory = [] - runtimes = [] - curr_job_id = futures[0].job_id - job_func = futures[0].function_name # each job is conducted on a single function + raise _missing_plotting_extra('job_summary') - for future in futures: - if curr_job_id != future.job_id: - cost = self.compute_handler.backend.calc_cost(runtimes, memory) - append([[curr_job_id, job_func, len(runtimes), sum(memory), - np.round(np.average(runtimes), 10), cost, ' ']]) + if not hasattr(self.compute_handler.backend, 'calc_cost'): + logger.warning( + f"Could not log job: {self.compute_handler.backend.name} " + "backend isn't supported by this function." + ) + return + + def append_rows(rows): + pd.DataFrame(rows).to_csv( + self.log_path, mode='a', header=False, index=False + ) - # updating next iteration's variables: - curr_job_id = future.job_id - job_func = future.function_name - memory.clear() - runtimes.clear() + if self.log_path: + # Carry over the cloud objects of the summary written last time, + # the last cell of the last row of its log + previous = pd.read_csv(self.log_path) + cloud_objects_n += float(previous.iloc[-1].iloc[-1]) + else: + self.log_path = os.path.join( + constants.LOGS_DIR, + datetime.now().strftime("%Y-%m-%d_%H-%M-%S.csv"), + ) - memory.append(future.runtime_memory) - runtimes.append(future.stats['worker_exec_time']) + # Writing the header alone overrides the summary of a previous call + headers = [ + 'Job_ID', 'Function', 'Invocations', 'Memory(MB)', + 'AvgRuntime', 'Cost', 'CloudObjects', + ] + pd.DataFrame([], columns=headers).to_csv(self.log_path, index=False) - # appends last Job-ID + futures = self._as_future_list(self.futures) + for job_id, job_func, runtimes, memory in _group_futures_by_job(futures): cost = self.compute_handler.backend.calc_cost(runtimes, memory) - append([[curr_job_id, job_func, len(runtimes), sum(memory), - np.round(np.average(runtimes), 10), cost, ' ']]) - # append summary row to end of the dataframe - append_summary() - - else: # calc_cost() doesn't exist for chosen computational backend. - logger.warning("Could not log job: {} backend isn't supported by this function." - .format(self.compute_handler.backend.name)) - return - logger.info("View log file logs at {}".format(self.log_path)) + append_rows([[ + job_id, job_func, len(runtimes), sum(memory), + np.round(np.average(runtimes), 10), cost, ' ', + ]]) + + summary = pd.read_csv(self.log_path) + total_average = ( + sum(summary.AvgRuntime * summary.Invocations) + / summary.Invocations.sum() + ) + append_rows([[ + 'Summary', + ' ', + summary.Invocations.sum(), + summary['Memory(MB)'].sum(), + round(total_average, 10), + summary.Cost.sum(), + cloud_objects_n, + ]]) + + logger.info(f"View log file logs at {self.log_path}") class LocalhostExecutor(FunctionExecutor): @@ -729,7 +1006,8 @@ class LocalhostExecutor(FunctionExecutor): :param storage: Name of the storage backend to use. :param monitoring: monitoring system. :param log_level: log level to use during the execution. - :param kwargs: Any parameter that can be set in the compute backend section of the config file, can be set here + :param kwargs: Any parameter that can be set in the compute + backend section of the config file, can be set here """ def __init__( @@ -738,8 +1016,8 @@ def __init__( config_file: Optional[str] = None, storage: Optional[str] = None, monitoring: Optional[str] = None, - log_level: Optional[str] = False, - **kwargs: Optional[Dict[str, Any]] + log_level: Union[str, bool, None] = False, + **kwargs: Any ): super().__init__( backend=LOCALHOST, @@ -752,18 +1030,10 @@ def __init__( ) -class ServerlessExecutor(FunctionExecutor): - """ - Initialize a ServerlessExecutor class. +class _FixedModeExecutor(FunctionExecutor): + """FunctionExecutor subclass that pins execution mode via `_mode`.""" - :param config: Settings passed in here will override those in config file - :param config_file: Path to the lithops config file - :param backend: Name of the serverless compute backend to use - :param storage: Name of the storage backend to use - :param monitoring: monitoring system - :param log_level: log level to use during the execution - :param kwargs: Any parameter that can be set in the compute backend section of the config file, can be set here - """ + _mode = None def __init__( self, @@ -772,13 +1042,13 @@ def __init__( backend: Optional[str] = None, storage: Optional[str] = None, monitoring: Optional[str] = None, - log_level: Optional[str] = False, - **kwargs: Optional[Dict[str, Any]] + log_level: Union[str, bool, None] = False, + **kwargs: Any ): super().__init__( config=config, config_file=config_file, - mode='serverless', + mode=self._mode, backend=backend, storage=storage, monitoring=monitoring, @@ -787,7 +1057,24 @@ def __init__( ) -class StandaloneExecutor(FunctionExecutor): +class ServerlessExecutor(_FixedModeExecutor): + """ + Initialize a ServerlessExecutor class. + + :param config: Settings passed in here will override those in config file + :param config_file: Path to the lithops config file + :param backend: Name of the serverless compute backend to use + :param storage: Name of the storage backend to use + :param monitoring: monitoring system + :param log_level: log level to use during the execution + :param kwargs: Any parameter that can be set in the compute + backend section of the config file, can be set here + """ + + _mode = SERVERLESS + + +class StandaloneExecutor(_FixedModeExecutor): """ Initialize a StandaloneExecutor class. @@ -799,23 +1086,4 @@ class StandaloneExecutor(FunctionExecutor): :param log_level: log level to use during the execution """ - def __init__( - self, - config: Optional[Dict[str, Any]] = None, - config_file: Optional[str] = None, - backend: Optional[str] = None, - storage: Optional[str] = None, - monitoring: Optional[str] = None, - log_level: Optional[str] = False, - **kwargs: Optional[Dict[str, Any]] - ): - super().__init__( - config=config, - config_file=config_file, - mode='standalone', - backend=backend, - storage=storage, - monitoring=monitoring, - log_level=log_level, - **kwargs, - ) + _mode = STANDALONE diff --git a/lithops/future.py b/lithops/future.py index cc465c445..17d71c2a5 100644 --- a/lithops/future.py +++ b/lithops/future.py @@ -33,16 +33,40 @@ create_job_key ) from lithops.constants import FN_LOG_FILE, LOGS_DIR +from lithops.utils import log_prefix logger = logging.getLogger(__name__) +_STAT_KEY_PREFIXES = ('func', 'host', 'worker') + + +def _stats_from_prefixed_keys(mapping: dict) -> dict: + """ + Picks the entries that hold a statistic, told apart from the rest of the + status by their key prefix + """ + return { + key: mapping[key] + for key in mapping + if any(key.startswith(p) for p in _STAT_KEY_PREFIXES) + } + + +def _pickle_from_encoded(encoded: str): + """ + Unpickles a value the worker put in the call status. It travels as the + repr() of its pickle, so eval() is what turns it back into bytes + """ + return pickle.loads(eval(encoded)) + class ResponseFuture: """ - Object representing the result of a Lithops invocation. Returns the status of the - execution and the result when available. + Result of a Lithops invocation. Exposes execution status + and the return value once it is available. """ - class State(): + + class State: New = "New" Invoked = "Invoked" Running = "Running" @@ -68,7 +92,7 @@ def __init__(self, call_id, job, job_metadata, storage_config): self._storage_config = storage_config self._produce_output = True self._read = False - self._state = ResponseFuture.State.New + self._state = self.State.New self._exception = Exception() self._handler_exception = False self._new_futures = None @@ -79,13 +103,16 @@ def __init__(self, call_id, job, job_metadata, storage_config): self._status_query_count = 0 self._output_query_count = 0 - for key in job_metadata: - if any(key.startswith(ss) for ss in ['func', 'host', 'worker']): - self.stats[key] = job_metadata[key] - + self.stats.update(_stats_from_prefixed_keys(job_metadata)) self._storage_path = get_storage_path(self._storage_config) - def _set_state(self, new_state): + def _id_prefix(self) -> str: + """ + Identity of the job this call belongs to, for the log messages + """ + return log_prefix(self.executor_id, self.job_id) + + def _set_state(self, new_state: str) -> None: self._state = new_state def cancel(self): @@ -96,260 +123,402 @@ def cancelled(self): @property def new(self): - return self._state == ResponseFuture.State.New + return self._state == self.State.New @property def invoked(self): - return self._state == ResponseFuture.State.Invoked + return self._state == self.State.Invoked @property def running(self): - return self._state == ResponseFuture.State.Running + return self._state == self.State.Running @property def ready(self): - return self._state == ResponseFuture.State.Ready + return self._state == self.State.Ready @property def error(self): - return self._state == ResponseFuture.State.Error + return self._state == self.State.Error @property def success(self): - return self._state in [ResponseFuture.State.Success, - ResponseFuture.State.Error] + return self._state in (self.State.Success, self.State.Error) @property def done(self): - return self._state in [ResponseFuture.State.Done, - ResponseFuture.State.Error, - ResponseFuture.State.Unknown] + return self._state in ( + self.State.Done, + self.State.Error, + self.State.Unknown, + ) @property def futures(self): return self._new_futures is not None def _set_invoked(self): - """ Set the future as invoked""" - self._state = ResponseFuture.State.Invoked + """Set the future as invoked""" + self._set_state(self.State.Invoked) def _set_running(self, call_status): - """ Set the future as running""" + """Set the future as running""" self._call_status = call_status self.activation_id = self._call_status['activation_id'] - self._state = ResponseFuture.State.Running + self._set_state(self.State.Running) def _set_exception(self): - """ Set the future as error""" + """Set the future as error""" self._read = True self._host_status_done_tstamp = time.time() if not self.done: - self._state = ResponseFuture.State.Unknown + self._set_state(self.State.Unknown) def _set_ready(self, call_status): - """ Set the future as ready""" + """Set the future as ready""" self._call_status = call_status self._host_status_done_tstamp = time.time() - self._state = ResponseFuture.State.Ready + self._set_state(self.State.Ready) def _set_futures(self, call_status): - """ Set the future as futures""" + """Set the future as futures""" self._call_status = call_status self._host_status_done_tstamp = time.time() self.status(throw_except=False) - self._state = ResponseFuture.State.Ready + self._set_state(self.State.Ready) def _set_mapreduce(self): - """ Set the future as mapreduce map""" + """Set the future as mapreduce map""" self._read = True self._produce_output = False if self.success: - self._state = ResponseFuture.State.Done + self._set_state(self.State.Done) - def status(self, throw_except=True, internal_storage=None, check_only=False, wait_dur_sec=1): + def _query_call_status(self, internal_storage): """ - Return the status returned by the call. - If the call raised an exception, this method will raise the same exception - If the future is cancelled before completing then CancelledError will be raised. + Reads the status the worker wrote, counting the query. Returns None + while the call has not finished + """ + status = internal_storage.get_call_status( + self.executor_id, self.job_id, self.call_id + ) + self._status_query_count += 1 + return status + + def _query_call_output(self, internal_storage): + """ + Reads the result the worker wrote, counting the query. Returns None + while it is not there yet + """ + output = internal_storage.get_call_output( + self.executor_id, self.job_id, self.call_id + ) + self._output_query_count += 1 + return output + + def _write_activation_logs(self) -> None: + """ + Replays the log of the activation, which travels compressed in the + status, into the job log and the global function log + """ + encoded = self._call_status['logs'].encode() + self.logs = zlib.decompress(base64.b64decode(encoded)).decode() + job_key = create_job_key(self.executor_id, self.job_id) + log_file = os.path.join(LOGS_DIR, job_key + '.log') + header = f"Activation: '{self.runtime_name}' ({self.activation_id})\n[\n" + # Every line is indented but the last one, so that the bracket that + # closes the activation stays at the left margin + newline_count = self.logs.count('\n') + indented = self.logs.replace('\r', '').replace( + '\n', '\n ', newline_count - 1 + ) + formatted = header + ' ' + indented + ']\n\n' + os.makedirs(LOGS_DIR, exist_ok=True) + for path in (log_file, FN_LOG_FILE): + with open(path, 'a') as lf: + lf.write(formatted) + + def _poll_until_ready(self, internal_storage, wait_dur_sec, check_only): + """ + Waits for the worker to write the status of the call, unless only + checking, in which case it returns whatever there is right away + """ + self._call_status = self._query_call_status(internal_storage) + if check_only: + return self._call_status + while self._call_status is None: + time.sleep(wait_dur_sec) + self._call_status = self._query_call_status(internal_storage) + self._host_status_done_tstamp = time.time() + return self._call_status + + def _raise_call_exception(self, throw_except): + """ + Rebuilds the exception the function raised and, unless the caller + asked not to, re-raises it with the traceback it had in the worker + """ + self._set_state(self.State.Error) + self._exception = _pickle_from_encoded( + self._call_status['exc_info'] + ) - :param check_only: Return None immediately if job is not complete. Default False. - :param throw_except: Reraise exception if call raised. Default true. - :param internal_storage: Storage handler to poll cloud storage. Default None. + if not self._call_status.get('exc_pickle_fail', False): + fn_exctype = self._exception[0] + fn_exc = self._exception[1] + # The worker marks its own failures with a HANDLER first argument. + # They carry no user traceback worth printing, so the marker is + # dropped and only the message is kept + if fn_exc.args and fn_exc.args[0] == "HANDLER": + self._handler_exception = True + try: + del fn_exc.errno + except Exception: + pass + fn_exc.args = (fn_exc.args[1],) + else: + fn_exctype = Exception + fn_exc = Exception(self._exception['exc_value']) + self._exception = ( + fn_exctype, + fn_exc, + self._exception['exc_traceback'], + ) + + logger.warning( + f'{self._id_prefix()} - CallID: {self.call_id} - ' + f'There was an exception - Activation ID: {self.activation_id} - {fn_exctype.__name__}' + ) + + # Reraising here would print a traceback pointing at this file, so the + # hook prints the one the function had in the worker instead. Anything + # else raised afterwards restores the default hook + def exception_hook(exctype, exc, trcbck): + if exctype == fn_exctype and str(exc) == str(fn_exc): + if self._handler_exception: + logger.warning( + f'Exception: {fn_exctype.__name__} - {fn_exc}' + ) + else: + traceback.print_exception(*self._exception) + else: + sys.excepthook = sys.__excepthook__ + traceback.print_exception(exctype, exc, trcbck) + + if throw_except: + sys.excepthook = exception_hook + reraise(*self._exception) + return None + + def _record_status_stats(self) -> float: + """ + Copies the statistics of the call status into the future and returns + how long the function ran for + """ + self.stats['host_status_done_tstamp'] = ( + self._host_status_done_tstamp or time.time() + ) + self.stats['host_status_query_count'] = self._status_query_count + self.stats.update(_stats_from_prefixed_keys(self._call_status)) + + exec_time = round( + self.stats['worker_end_tstamp'] + - self.stats['worker_start_tstamp'], + 8, + ) + self.stats['worker_exec_time'] = exec_time + return exec_time + + def _resolve_new_futures(self) -> None: + """ + Adopts the futures the function returned: this call produces no + result of its own, the client has to wait for those instead + """ + new_futures = _pickle_from_encoded( + self._call_status['new_futures'] + ) + if isinstance(new_futures, ResponseFuture): + self._new_futures = [new_futures] + else: + self._new_futures = new_futures + + def _read_inline_result(self) -> None: + """ + Takes the result the worker embedded in the status, which saves the + client one storage request + """ + self._call_output = _pickle_from_encoded( + self._call_status['result'] + ) + self.stats['host_result_done_tstamp'] = time.time() + self.stats['host_result_query_count'] = 0 + logger.debug( + f'{self._id_prefix()} - Got output from call ' + f'{self.call_id} - Activation ID: {self.activation_id}' + ) + + def status( + self, + throw_except=True, + internal_storage=None, + check_only=False, + wait_dur_sec=1, + ): + """ + Return the status returned by the call. + If the call raised an exception, this method will raise + the same exception. If the future is cancelled before + completing then CancelledError will be raised. + + :param check_only: Return None immediately if job is + not complete. Default False. + :param throw_except: Reraise exception if call raised. + Default true. + :param internal_storage: Storage handler to poll cloud + storage. Default None. :param wait_dur_sec: Time interval between each check :return: Result of the call. - :raises CancelledError: If the job is cancelled before completed. - :raises TimeoutError: If job is not complete after `timeout` seconds. + :raises CancelledError: If the job is cancelled + before completed. + :raises TimeoutError: If job is not complete after + `timeout` seconds. """ - if self._state == ResponseFuture.State.New: + if self._state == self.State.New: raise ValueError("task not yet invoked") if self.success or self.done: return self._call_status - if self._call_status is None or self._call_status['type'] == '__init__': + needs_fetch = ( + self._call_status is None + or self._call_status['type'] == '__init__' + ) + if needs_fetch: if internal_storage is None: internal_storage = InternalStorage(self._storage_config) - check_storage_path(internal_storage.get_storage_config(), self._storage_path) - self._call_status = internal_storage.get_call_status(self.executor_id, self.job_id, self.call_id) - self._status_query_count += 1 - + check_storage_path( + internal_storage.get_storage_config(), + self._storage_path, + ) + status = self._poll_until_ready( + internal_storage, wait_dur_sec, check_only + ) if check_only: - return self._call_status - - while self._call_status is None: - time.sleep(wait_dur_sec) - self._call_status = internal_storage.get_call_status(self.executor_id, self.job_id, self.call_id) - self._status_query_count += 1 - self._host_status_done_tstamp = time.time() + return status - self.stats['host_status_done_tstamp'] = self._host_status_done_tstamp or time.time() - self.stats['host_status_query_count'] = self._status_query_count self.activation_id = self._call_status['activation_id'] if 'logs' in self._call_status: - self.logs = zlib.decompress(base64.b64decode(self._call_status['logs'].encode())).decode() - job_key = create_job_key(self.executor_id, self.job_id) - log_file = os.path.join(LOGS_DIR, job_key + '.log') - header = "Activation: '{}' ({})\n[\n".format(self.runtime_name, self.activation_id) - tail = ']\n\n' - output = self.logs.replace('\r', '').replace('\n', '\n ', self.logs.count('\n') - 1) - with open(log_file, 'a') as lf: - lf.write(header + ' ' + output + tail) - with open(FN_LOG_FILE, 'a') as lf: - lf.write(header + ' ' + output + tail) - - for key in self._call_status: - if any(key.startswith(ss) for ss in ['func', 'host', 'worker']): - self.stats[key] = self._call_status[key] - - self.stats['worker_exec_time'] = round(self.stats['worker_end_tstamp'] - self.stats['worker_start_tstamp'], 8) - total_time = format(round(self.stats['worker_exec_time'], 2), '.2f') + self._write_activation_logs() + + exec_time = self._record_status_stats() logger.debug( - f'ExecutorID {self.executor_id} | JobID {self.job_id} - Got status from call {self.call_id} ' - f'- Activation ID: {self.activation_id} - Time: {str(total_time)} seconds' + f'{self._id_prefix()} - Got status from call ' + f'{self.call_id} - Activation ID: {self.activation_id} ' + f'- Time: {exec_time:.2f} seconds' ) if self._call_status['exception']: - self._set_state(ResponseFuture.State.Error) - self._exception = pickle.loads(eval(self._call_status['exc_info'])) - - if not self._call_status.get('exc_pickle_fail', False): - fn_exctype = self._exception[0] - fn_exc = self._exception[1] - if fn_exc.args and fn_exc.args[0] == "HANDLER": - self._handler_exception = True - try: - del fn_exc.errno - except Exception: - pass - fn_exc.args = (fn_exc.args[1],) - else: - fn_exctype = Exception - fn_exc = Exception(self._exception['exc_value']) - self._exception = (fn_exctype, fn_exc, - self._exception['exc_traceback']) - - logger.warning( - 'ExecutorID {} | JobID {} - CallID: {} - There was an exception - Activation ID: {} - {}' - .format(self.executor_id, self.job_id, self.call_id, self.activation_id, fn_exctype.__name__) - ) - - def exception_hook(exctype, exc, trcbck): - if exctype == fn_exctype and str(exc) == str(fn_exc): - if self._handler_exception: - logger.warning(f'Exception: {fn_exctype.__name__} - {fn_exc}') - else: - traceback.print_exception(*self._exception) - else: - sys.excepthook = sys.__excepthook__ - traceback.print_exception(exctype, exc, trcbck) - - if throw_except: - sys.excepthook = exception_hook - reraise(*self._exception) - else: - return None + return self._raise_call_exception(throw_except) if 'new_futures' in self._call_status and not self._new_futures: - new_futures = pickle.loads(eval(self._call_status['new_futures'])) - self._new_futures = [new_futures] if type(new_futures) is ResponseFuture else new_futures - + self._resolve_new_futures() elif self._call_status['func_result_size'] == 0: self._produce_output = False if 'result' in self._call_status: - self._call_output = pickle.loads(eval(self._call_status['result'])) - self.stats['host_result_done_tstamp'] = time.time() - self.stats['host_result_query_count'] = 0 - logger.debug( - f'ExecutorID {self.executor_id} | JobID {self.job_id} - Got output ' - f'from call {self.call_id} - Activation ID: {self.activation_id}' - ) + self._read_inline_result() if self._call_output is not None or not self._produce_output: - self._set_state(ResponseFuture.State.Done) + self._set_state(self.State.Done) else: - self._set_state(ResponseFuture.State.Success) + self._set_state(self.State.Success) return self._call_status - def result(self, throw_except=True, internal_storage=None, retries=10, wait_dur_sec=1): + def _fetch_call_output(self, internal_storage, retries, wait_dur_sec): + """ + Reads the result of the call from the storage, retrying while it is + not there. The status can arrive before the result. Returns None if + it never shows up + """ + call_output = self._query_call_output(internal_storage) + while ( + call_output is None + and self._output_query_count < retries + ): + time.sleep(wait_dur_sec) + call_output = self._query_call_output(internal_storage) + return call_output + + def result( + self, + throw_except=True, + internal_storage=None, + retries=10, + wait_dur_sec=1, + ): """ Return the value returned by the call. - If the call raised an exception, this method will raise the same exception - If the future is cancelled before completing then CancelledError will be raised. - - :param throw_except: Reraise exception if call raised. Default true. - :param internal_storage: Storage handler to poll cloud storage. Default None. - :param retries: Number of times to check if the result file is in the storage - :param wait_dur_sec: Time interval between each retry check + If the call raised an exception, this method will raise + the same exception. If the future is cancelled before + completing then CancelledError will be raised. + + :param throw_except: Reraise exception if call raised. + Default true. + :param internal_storage: Storage handler to poll cloud + storage. Default None. + :param retries: Number of times to check if the result + file is in the storage + :param wait_dur_sec: Time interval between each retry :return: Result of the call. - :raises CancelledError: If the job is cancelled before completed. - :raises TimeoutError: If job is not complete after `timeout` seconds. + :raises CancelledError: If the job is cancelled + before completed. + :raises TimeoutError: If job is not complete after + `timeout` seconds. """ - if self._state == ResponseFuture.State.New: + if self._state == self.State.New: raise ValueError("Task not yet invoked") if not self.done and internal_storage is None: - internal_storage = InternalStorage(storage_config=self._storage_config) + internal_storage = InternalStorage(self._storage_config) - self.status(throw_except=throw_except, internal_storage=internal_storage, wait_dur_sec=wait_dur_sec) + self.status( + throw_except=throw_except, + internal_storage=internal_storage, + wait_dur_sec=wait_dur_sec, + ) if self.futures: self._call_output = self._new_futures - self._set_state(ResponseFuture.State.Done) + self._set_state(self.State.Done) if self.done: return self._call_output if self._call_output is None: - call_output = internal_storage.get_call_output(self.executor_id, self.job_id, self.call_id) - self._output_query_count += 1 - - while call_output is None and self._output_query_count < retries: - time.sleep(wait_dur_sec) - call_output = internal_storage.get_call_output(self.executor_id, self.job_id, self.call_id) - self._output_query_count += 1 + call_output = self._fetch_call_output( + internal_storage, retries, wait_dur_sec + ) if call_output is None: if throw_except: raise Exception( - f'ExecutorID {self.executor_id} | JobID {self.job_id} - Unable to get ' + f'{self._id_prefix()} - Unable to get ' f'the result from call {self.call_id} - Activation ID: {self.activation_id}' ) - else: - self._set_state(ResponseFuture.State.Error) - return None + self._set_state(self.State.Error) + return None self._call_output = pickle.loads(call_output) - self.stats['host_result_done_tstamp'] = time.time() self.stats['host_result_query_count'] = self._output_query_count - logger.debug(f'ExecutorID {self.executor_id} | JobID {self.job_id} - Got output ' - f'from call {self.call_id} - Activation ID: {self.activation_id}') + logger.debug( + f'{self._id_prefix()} - Got output from call ' + f'{self.call_id} - Activation ID: {self.activation_id}' + ) - self._set_state(ResponseFuture.State.Done) + self._set_state(self.State.Done) return self._call_output diff --git a/lithops/invokers.py b/lithops/invokers.py index 607308313..78fefdc00 100644 --- a/lithops/invokers.py +++ b/lithops/invokers.py @@ -22,6 +22,7 @@ import shutil import logging import threading +from math import ceil from concurrent.futures import ThreadPoolExecutor from lithops.future import ResponseFuture @@ -32,7 +33,9 @@ version_str, is_lithops_worker, iterchunks, - BackendType + monitoring_queues, + BackendType, + log_prefix, ) from lithops.constants import ( LOGGER_LEVEL, @@ -46,27 +49,74 @@ logger = logging.getLogger(__name__) -def create_invoker(config, executor_id, internal_storage, - compute_handler, job_monitor): +def create_invoker( + config, + executor_id, + internal_storage, + compute_handler, + job_monitor, +): """ Creates the appropriate invoker based on the backend type """ - if compute_handler.get_backend_type() == BackendType.BATCH.value: - return BatchInvoker( - config, - executor_id, - internal_storage, - compute_handler, - job_monitor + invoker_cls = { + BackendType.BATCH.value: BatchInvoker, + BackendType.FAAS.value: FaaSInvoker, + }.get(compute_handler.get_backend_type()) + if invoker_cls is None: + return None + return invoker_cls( + config, + executor_id, + internal_storage, + compute_handler, + job_monitor, + ) + + +def _format_call_id(index: int) -> str: + """ + Formats a call index as the fixed width call id the workers expect + """ + return f'{index:05d}' + + +def _timed_invoke(compute_handler, payload): + """ + Invokes the payload and returns the activation id together with how long + the backend took to accept it, already formatted for the log + """ + start = time.time() + activation_id = compute_handler.invoke(payload) + return activation_id, f'{round(time.time() - start, 3):.3f}' + + +def _raise_invoke_error(invoke_future) -> None: + """ + Done callback of an invocation nobody waits for. Re-raising here makes + concurrent.futures log the failure instead of dropping it silently + """ + invoke_future.result() + + +def _verify_runtime_meta(runtime_meta, runtime_name): + """ + Ensures the runtime runs the same Lithops and Python versions as this + client, as it has to unpickle the function this client serializes + """ + if __version__ != runtime_meta['lithops_version']: + raise Exception( + f"Lithops version mismatch. Host version: {__version__} - " + f"Runtime version: {runtime_meta['lithops_version']}" ) - elif compute_handler.get_backend_type() == BackendType.FAAS.value: - return FaaSInvoker( - config, - executor_id, - internal_storage, - compute_handler, - job_monitor + py_local_version = version_str(sys.version_info) + py_remote_version = runtime_meta['python_version'] + if py_local_version != py_remote_version: + raise Exception( + f"The indicated runtime '{runtime_name}' is running Python " + f"{py_remote_version} and it is not compatible with the local " + f"Python version {py_local_version}" ) @@ -75,10 +125,19 @@ class Invoker: Abstract invoker class """ - def __init__(self, config, executor_id, internal_storage, compute_handler, job_monitor): + def __init__( + self, + config, + executor_id, + internal_storage, + compute_handler, + job_monitor, + ): log_level = logger.getEffectiveLevel() self.log_active = log_level != logging.WARNING - self.log_level = LOGGER_LEVEL if not self.log_active else log_level + self.log_level = ( + LOGGER_LEVEL if not self.log_active else log_level + ) self.config = config self.executor_id = executor_id @@ -94,7 +153,9 @@ def __init__(self, config, executor_id, internal_storage, compute_handler, job_m self.mode = self.config['lithops']['mode'] self.backend = self.config['lithops']['backend'] - self.include_function = self.config[self.backend].get('runtime_include_function', False) + self.include_function = self.config[self.backend].get( + 'runtime_include_function', False + ) self.runtime_info = self.compute_handler.get_runtime_info() self.runtime_name = self.runtime_info['runtime_name'] @@ -102,52 +163,67 @@ def __init__(self, config, executor_id, internal_storage, compute_handler, job_m verify_runtime_name(self.runtime_name) - logger.debug(f'ExecutorID {self.executor_id} - Invoker initialized.' - f' Max workers: {self.max_workers}') + logger.debug( + f'{log_prefix(self.executor_id)} - Invoker initialized. Max workers: {self.max_workers}' + ) + + def _deploy_runtime(self, runtime_key, runtime_memory): + """ + Deploys the selected runtime and caches its metadata, so that the + next job that selects it finds it already deployed + """ + msg = f'Runtime {self.runtime_name}' + if runtime_memory: + msg += f' with {runtime_memory}MB' + logger.info(f'{msg} is not yet deployed') + + runtime_timeout = self.runtime_info['runtime_timeout'] + runtime_meta = self.compute_handler.deploy_runtime( + self.runtime_name, + runtime_memory, + runtime_timeout, + ) + runtime_meta['runtime_timeout'] = runtime_timeout + self.internal_storage.put_runtime_meta(runtime_key, runtime_meta) + return runtime_meta def select_runtime(self, job_id, runtime_memory): """ Return the runtime metadata """ - runtime_memory = runtime_memory or self.runtime_info['runtime_memory'] \ - if self.mode == SERVERLESS else self.runtime_info['runtime_memory'] - runtime_timeout = self.runtime_info['runtime_timeout'] + default_memory = self.runtime_info['runtime_memory'] + runtime_memory = ( + runtime_memory or default_memory + if self.mode == SERVERLESS + else default_memory + ) - msg = ('ExecutorID {} | JobID {} - Selected Runtime: {} ' - .format(self.executor_id, job_id, self.runtime_name)) - msg = msg + f'- {runtime_memory}MB' if runtime_memory else msg + msg = ( + f'{log_prefix(self.executor_id, job_id)} - ' + f'Selected Runtime: {self.runtime_name} ' + ) + if runtime_memory: + msg += f'- {runtime_memory}MB' logger.info(msg) - runtime_key = self.compute_handler.get_runtime_key(self.runtime_name, runtime_memory, __version__) - runtime_meta = self.internal_storage.get_runtime_meta(runtime_key) + runtime_key = self.compute_handler.get_runtime_key( + self.runtime_name, runtime_memory, __version__ + ) + runtime_meta = self.internal_storage.get_runtime_meta( + runtime_key + ) if not runtime_meta: - msg = f'Runtime {self.runtime_name}' - msg = msg + f' with {runtime_memory}MB' if runtime_memory else msg - logger.info(msg + ' is not yet deployed') - runtime_meta = self.compute_handler.deploy_runtime(self.runtime_name, runtime_memory, runtime_timeout) - runtime_meta['runtime_timeout'] = runtime_timeout - self.internal_storage.put_runtime_meta(runtime_key, runtime_meta) - - # Verify python version and lithops version - if __version__ != runtime_meta['lithops_version']: - raise Exception("Lithops version mismatch. Host version: {} - Runtime version: {}" - .format(__version__, runtime_meta['lithops_version'])) - - py_local_version = version_str(sys.version_info) - py_remote_version = runtime_meta['python_version'] - if py_local_version != py_remote_version: - raise Exception(("The indicated runtime '{}' is running Python {} and it " - "is not compatible with the local Python version {}") - .format(self.runtime_name, py_remote_version, py_local_version)) + runtime_meta = self._deploy_runtime(runtime_key, runtime_memory) + _verify_runtime_meta(runtime_meta, self.runtime_name) return runtime_meta def _create_payload(self, job): """ Creates the default payload dictionary """ - payload = { + return { 'config': self.config, 'chunksize': job.chunksize, 'log_level': self.log_level, @@ -159,6 +235,7 @@ def _create_payload(self, job): 'execution_timeout': job.execution_timeout, 'data_byte_ranges': job.data_byte_ranges, 'executor_id': job.executor_id, + 'monitoring_queues': monitoring_queues(job.executor_id), 'job_id': job.job_id, 'job_key': job.job_key, 'max_workers': self.max_workers, @@ -170,77 +247,85 @@ def _create_payload(self, job): 'worker_processes': job.worker_processes } - return payload - - def _run_job(self, job): + def _send_job_metrics(self, job): """ - Run a job + Reports the size of the job to Prometheus, if telemetry is enabled """ - if self.include_function: - logger.debug('ExecutorID {} | JobID {} - Runtime include function feature ' - ' is activated' .format(job.executor_id, job.job_id)) - job.runtime_name = self.runtime_name - extend_runtime(job, self.compute_handler, self.internal_storage) - self.runtime_name = job.runtime_name - - logger.info( - f'ExecutorID {job.executor_id} | JobID {job.job_id} - Starting function ' - f'invocation: {job.function_name}() - Total: {job.total_calls} activations' + labels = ( + ('job_id', job.job_key), + ('function_name', job.function_name), ) - self.prometheus.send_metric( name='job_total_calls', value=job.total_calls, type='counter', - labels=( - ('job_id', job.job_key), - ('function_name', job.function_name) - ) + labels=labels, ) - self.prometheus.send_metric( name='job_runtime_memory', value=job.runtime_memory or 0, type='counter', - labels=( - ('job_id', job.job_key), - ('function_name', job.function_name) + labels=labels, + ) + + def _build_futures(self, job): + """ + Creates one future per call of the job, already marked as invoked + """ + futures = [] + for i in range(job.total_calls): + fut = ResponseFuture( + _format_call_id(i), + job, + job.metadata.copy(), + self.storage_config, + ) + fut._set_state(ResponseFuture.State.Invoked) + futures.append(fut) + job.futures = futures + return futures + + def _run_job(self, job): + """ + Invokes a job through the backend specific _invoke_job() and returns + its futures. Stops the invoker if the invocation fails halfway + """ + prefix = log_prefix(job.executor_id, job.job_id) + if self.include_function: + logger.debug( + f'{prefix} - Runtime include function feature is activated' ) + job.runtime_name = self.runtime_name + extend_runtime( + job, self.compute_handler, self.internal_storage + ) + self.runtime_name = job.runtime_name + + logger.info( + f'{prefix} - Starting function invocation: {job.function_name}() - Total: ' + f'{job.total_calls} activations' ) + self._send_job_metrics(job) + if self.backend not in STANDALONE_BACKENDS: logger.debug( - f'ExecutorID {job.executor_id} | JobID {job.job_id} - Worker processes: ' + f'{prefix} - Worker processes: ' f'{job.worker_processes} - Chunksize: {job.chunksize}' ) try: job.runtime_name = self.runtime_name self._invoke_job(job) - except (KeyboardInterrupt, Exception) as e: + except (KeyboardInterrupt, Exception): self.stop() - raise e + raise log_file = os.path.join(LOGS_DIR, job.job_key + '.log') - logger.info( - f'ExecutorID {job.executor_id} | JobID {job.job_id} - View execution logs at {log_file}' - ) - - # Create all futures - futures = [] - for i in range(job.total_calls): - call_id = "{:05d}".format(i) - fut = ResponseFuture(call_id, job, - job.metadata.copy(), - self.storage_config) - fut._set_state(ResponseFuture.State.Invoked) - futures.append(fut) - - job.futures = futures - - return futures + logger.info(f'{prefix} - View execution logs at {log_file}') + return self._build_futures(job) - def stop(self): + def stop(self, wait: bool = False): """ Stop invoker-related processes """ @@ -249,27 +334,42 @@ def stop(self): class BatchInvoker(Invoker): """ - Module responsible to perform the invocations against a batch backend + Module responsible to perform the invocations against a + batch backend """ - def __init__(self, config, executor_id, internal_storage, compute_handler, job_monitor): - super().__init__(config, executor_id, internal_storage, compute_handler, job_monitor) + def __init__( + self, + config, + executor_id, + internal_storage, + compute_handler, + job_monitor, + ): + super().__init__( + config, + executor_id, + internal_storage, + compute_handler, + job_monitor, + ) self.compute_handler.init() def _invoke_job(self, job): """ - Run a job + Invokes every call of the job in a single request, as a batch backend + schedules the calls itself """ payload = self._create_payload(job) - payload['call_ids'] = ["{:05d}".format(i) for i in range(job.total_calls)] - - start = time.time() - activation_id = self.compute_handler.invoke(payload) - roundtrip = time.time() - start - resp_time = format(round(roundtrip, 3), '.3f') + payload['call_ids'] = [ + _format_call_id(i) for i in range(job.total_calls) + ] + activation_id, resp_time = _timed_invoke( + self.compute_handler, payload + ) logger.debug( - f'ExecutorID {job.executor_id} | JobID {job.job_id} - Job invoked ' + f'{log_prefix(job.executor_id, job.job_id)} - Job invoked ' f'({resp_time}s) - Activation ID: {activation_id or job.job_key}' ) @@ -279,115 +379,162 @@ def run_job(self, job): """ futures = self._run_job(job) self.job_monitor.start(futures) - return futures class FaaSInvoker(Invoker): """ - Module responsible to perform the invocations against a FaaS backend + Module responsible to perform the invocations against a + FaaS backend """ ASYNC_INVOKERS = 2 + # Upper bound for the wait of stop(wait=True) on each async invoker + STOP_TIMEOUT = 10 + + def __init__( + self, + config, + executor_id, + internal_storage, + compute_handler, + job_monitor, + ): + super().__init__( + config, + executor_id, + internal_storage, + compute_handler, + job_monitor, + ) - def __init__(self, config, executor_id, internal_storage, compute_handler, job_monitor): - super().__init__(config, executor_id, internal_storage, compute_handler, job_monitor) - - remote_invoker = self.config[self.backend].get('remote_invoker', False) - self.remote_invoker = remote_invoker if not is_lithops_worker() else False + remote_invoker = self.config[self.backend].get( + 'remote_invoker', False + ) + self.remote_invoker = ( + remote_invoker if not is_lithops_worker() else False + ) self.invokers = [] - self.ongoing_activations = 0 self.pending_calls_q = queue.Queue() self.should_run = False + self.running_workers = 0 self.sync = is_lithops_worker() - self.invoke_pool_threads = self.config[self.backend]['invoke_pool_threads'] + self.invoke_pool_threads = self.config[self.backend][ + 'invoke_pool_threads' + ] self.executor = ThreadPoolExecutor(self.invoke_pool_threads) - logger.debug(f'ExecutorID {self.executor_id} - Serverless invoker created') + logger.debug( + f'{log_prefix(self.executor_id)} - Serverless invoker created' + ) - def _start_async_invokers(self): - """Starts the invoker process responsible to spawn pending calls - in background. + def _async_invoker_loop(self, inv_id): """ + Token bucket scheduling loop: spends one token, which the monitor + puts for every worker that becomes free, on the next pending chunk + of calls. Runs in a background thread until stop() is called + """ + logger.debug( + f'{log_prefix(self.executor_id)} - Async invoker {inv_id} started' + ) + workers = min(64, self.invoke_pool_threads // 4) + with ThreadPoolExecutor(max_workers=workers) as pool: + while self.should_run: + try: + self.job_monitor.token_bucket_q.get() + job, call_ids_range = self.pending_calls_q.get() + except KeyboardInterrupt: + break + if not self.should_run: + break + pool.submit(self._invoke_task, job, call_ids_range) - def invoker_process(inv_id): - """Run process that implements token bucket scheduling approach""" - logger.debug(f'ExecutorID {self.executor_id} - Async invoker {inv_id} started') - - with ThreadPoolExecutor(max_workers=min(64, self.invoke_pool_threads // 4)) as executor: - while self.should_run: - try: - self.job_monitor.token_bucket_q.get() - job, call_ids_range = self.pending_calls_q.get() - except KeyboardInterrupt: - break - if self.should_run: - executor.submit(self._invoke_task, job, call_ids_range) - else: - break - - logger.debug(f'ExecutorID {self.executor_id} - Async invoker {inv_id} finished') + logger.debug( + f'{log_prefix(self.executor_id)} - Async invoker {inv_id} finished' + ) + def _start_async_invokers(self): + """Starts the invoker process responsible to spawn + pending calls in background. + """ for inv_id in range(self.ASYNC_INVOKERS): self.job_monitor.token_bucket_q.put('#') - p = threading.Thread(target=invoker_process, args=(inv_id,)) - self.invokers.append(p) - p.daemon = True - p.start() + invoker = threading.Thread( + target=self._async_invoker_loop, args=(inv_id,) + ) + self.invokers.append(invoker) + invoker.daemon = True + invoker.start() - def stop(self): + def stop(self, wait: bool = False): """ - Stop async invokers + Stop async invokers. With wait, also waits for the threads to exit, + which they only do once the invocations already in flight are done """ if self.invokers: - logger.debug(f'ExecutorID {self.executor_id} - Stopping async invokers') + logger.debug( + f'{log_prefix(self.executor_id)} - Stopping async invokers' + ) self.should_run = False while not self.pending_calls_q.empty(): try: - self.pending_calls_q.get(False) - except Exception: - pass + self.pending_calls_q.get(block=False) + except queue.Empty: + break - for invoker in self.invokers: + # One sentinel per invoker, each one preceded by the token it + # blocks on, so that every loop wakes up and sees should_run + for _ in self.invokers: self.job_monitor.token_bucket_q.put('$') self.pending_calls_q.put((None, None)) - self.invokers = [] + invokers, self.invokers = self.invokers, [] + + if wait: + # The loops leave their thread pool behind, and it only drains + # the invocations already in flight once they exit, so callers + # that cannot outlive that have to wait for it + current_thread = threading.current_thread() + for invoker in invokers: + if invoker is not current_thread: + invoker.join(timeout=self.STOP_TIMEOUT) def _invoke_task(self, job, call_ids_range): - """Method used to perform the actual invocation against the - compute backend. """ - # prepare payload + Invokes one chunk of calls against the compute backend. A backend + that refuses the invocation returns no activation id, in which case + the chunk goes back to the pending queue with its token + """ payload = self._create_payload(job) - - call_ids = ["{:05d}".format(i) for i in call_ids_range] + call_ids = [_format_call_id(i) for i in call_ids_range] payload['call_ids'] = call_ids if job.data_key: - data_byte_ranges = [job.data_byte_ranges[int(call_id)] for call_id in call_ids] - payload['data_byte_ranges'] = data_byte_ranges + payload['data_byte_ranges'] = [ + job.data_byte_ranges[int(call_id)] + for call_id in call_ids + ] else: del payload['data_byte_ranges'] - payload['data_byte_strs'] = [job.data_byte_strs[int(call_id)] for call_id in call_ids] + payload['data_byte_strs'] = [ + job.data_byte_strs[int(call_id)] + for call_id in call_ids + ] - # do the invocation - start = time.time() - activation_id = self.compute_handler.invoke(payload) - roundtrip = time.time() - start - resp_time = format(round(roundtrip, 3), '.3f') + activation_id, resp_time = _timed_invoke( + self.compute_handler, payload + ) if not activation_id: - # reached quota limit time.sleep(random.randint(0, 5)) self.pending_calls_q.put((job, call_ids_range)) self.job_monitor.token_bucket_q.put('#') return logger.debug( - f'ExecutorID {job.executor_id} | JobID {job.job_id} - Calls {", ".join(call_ids)} ' + f'{log_prefix(job.executor_id, job.job_id)} - Calls {", ".join(call_ids)} ' f'invoked ({resp_time}s) - Activation ID: {activation_id}' ) @@ -395,26 +542,67 @@ def _invoke_job_remote(self, job): """ Logic for invoking a job using a remote function """ - start = time.time() - payload = {} - payload['config'] = self.config - payload['log_level'] = self.log_level - payload['runtime_name'] = job.runtime_name - payload['runtime_memory'] = job.runtime_memory - payload['remote_invoker'] = True - payload['job'] = job.__dict__ - - activation_id = self.compute_handler.invoke(payload) - roundtrip = time.time() - start - resp_time = format(round(roundtrip, 3), '.3f') + payload = { + 'config': self.config, + 'log_level': self.log_level, + 'runtime_name': job.runtime_name, + 'runtime_memory': job.runtime_memory, + 'remote_invoker': True, + 'job': job.__dict__, + } + activation_id, resp_time = _timed_invoke( + self.compute_handler, payload + ) if activation_id: logger.debug( - f'ExecutorID {job.executor_id} | JobID {job.job_id} - Remote invoker ' + f'{log_prefix(job.executor_id, job.job_id)} - Remote invoker ' f'call done ({resp_time}s) - Activation ID: {activation_id}' ) - else: - raise Exception('Unable to spawn remote invoker') + return + raise Exception('Unable to spawn remote invoker') + + def _drain_token_bucket(self): + """ + Takes back the tokens left over by previous jobs, one per worker that + already finished, so that this job can reuse those workers + """ + if self.running_workers <= 0: + return + + while not self.job_monitor.token_bucket_q.empty(): + try: + self.job_monitor.token_bucket_q.get(block=False) + except queue.Empty: + break + self.running_workers -= 1 + if self.running_workers == 0: + break + + def _queue_call_ranges(self, job, call_ids): + """ + Leaves the calls in the pending queue, in chunks of one worker each, + for the async invokers to pick up as tokens become available + """ + for call_ids_range in iterchunks(call_ids, job.chunksize): + self.pending_calls_q.put((job, call_ids_range)) + + def _invoke_direct(self, job, call_ids): + """ + Invokes the given calls right away, one worker per chunk. Inside a + worker there is no async invoker, so it waits for them to be invoked + """ + invoke_futures = [] + for call_ids_range in iterchunks(call_ids, job.chunksize): + invoke_future = self.executor.submit( + self._invoke_task, job, call_ids_range + ) + invoke_future.add_done_callback(_raise_invoke_error) + invoke_futures.append(invoke_future) + + if self.sync: + for invoke_future in invoke_futures: + invoke_future.result() def _invoke_job(self, job): """ @@ -426,66 +614,46 @@ def _invoke_job(self, job): if self.remote_invoker: return self._invoke_job_remote(job) - if self.should_run is False: + prefix = log_prefix(job.executor_id, job.job_id) + + if not self.should_run: self.running_workers = 0 self.should_run = True self._start_async_invokers() - if self.running_workers > 0 and not self.job_monitor.token_bucket_q.empty(): - while not self.job_monitor.token_bucket_q.empty(): - try: - self.job_monitor.token_bucket_q.get(False) - self.running_workers -= 1 - if self.running_workers == 0: - break - except Exception: - pass - - if self.running_workers < self.max_workers: - free_workers = self.max_workers - self.running_workers - total_direct = free_workers * job.chunksize - callids = range(job.total_calls) - callids_to_invoke_direct = callids[:total_direct] - callids_to_invoke_nondirect = callids[total_direct:] - - ci = len(callids_to_invoke_direct) - cz = job.chunksize - consumed_workers = ci // cz + (ci % cz > 0) - self.running_workers += consumed_workers + self._drain_token_bucket() + if self.running_workers >= self.max_workers: logger.debug( - f'ExecutorID {job.executor_id} | JobID {job.job_id} - Free workers: ' - f'{free_workers} - Going to run {len(callids_to_invoke_direct)} activations ' - f'in {consumed_workers} workers' + f'{prefix} - Reached maximum ' + f'{self.max_workers} workers, queuing {job.total_calls} ' + f'function activations' ) + self._queue_call_ranges(job, range(job.total_calls)) + return - def _callback(future): - future.result() - - invoke_futures = [] - for call_ids_range in iterchunks(callids_to_invoke_direct, job.chunksize): - future = self.executor.submit(self._invoke_task, job, call_ids_range) - future.add_done_callback(_callback) - invoke_futures.append(future) - - if self.sync: - [f.result() for f in invoke_futures] - - # Put into the queue the rest of the callids to invoke within the process - if callids_to_invoke_nondirect: - logger.debug( - f'ExecutorID {job.executor_id} | JobID {job.job_id} - Putting remaining ' - f'{len(callids_to_invoke_nondirect)} function activations into pending queue' - ) - for call_ids_range in iterchunks(callids_to_invoke_nondirect, job.chunksize): - self.pending_calls_q.put((job, call_ids_range)) - else: + free_workers = self.max_workers - self.running_workers + call_ids = range(job.total_calls) + direct = call_ids[:free_workers * job.chunksize] + queued = call_ids[free_workers * job.chunksize:] + + # One worker runs one chunk of calls, and the last one may be partial + consumed_workers = ceil(len(direct) / job.chunksize) + self.running_workers += consumed_workers + + logger.debug( + f'{prefix} - Free workers: ' + f'{free_workers} - Going to run {len(direct)} ' + f'activations in {consumed_workers} workers' + ) + self._invoke_direct(job, direct) + + if queued: logger.debug( - f'ExecutorID {job.executor_id} | JobID {job.job_id} - Reached maximum {self.max_workers} ' - f'workers, queuing {job.total_calls} function activations' + f'{prefix} - Putting remaining ' + f'{len(queued)} function activations into pending queue' ) - for call_ids_range in iterchunks(range(job.total_calls), job.chunksize): - self.pending_calls_q.put((job, call_ids_range)) + self._queue_call_ranges(job, queued) def run_job(self, job): """ @@ -498,58 +666,56 @@ def run_job(self, job): chunksize=job.chunksize, generate_tokens=True ) - return futures -def extend_runtime(job, compute_handler, internal_storage): +def _build_extended_runtime(job, compute_handler, base_docker_image): """ - This method is used when runtime_include_function is active + Builds an image that adds the function and its modules on top of the base + one. The build runs from the temporary directory holding them, as its + contents are the build context, and it is removed afterwards """ + ext_docker_file = '/'.join([job.local_tmp_dir, "Dockerfile"]) + with open(ext_docker_file, 'w') as df: + df.write('\n'.join([ + f'FROM {base_docker_image}', + f'ENV PYTHONPATH={SA_INSTALL_DIR}/modules:$PYTHONPATH', + f'COPY . {SA_INSTALL_DIR}' + ])) + + cwd = os.getcwd() + os.chdir(job.local_tmp_dir) + try: + compute_handler.build_runtime(job.runtime_name, ext_docker_file) + finally: + os.chdir(cwd) + shutil.rmtree(job.local_tmp_dir, ignore_errors=True) - base_docker_image = job.runtime_name - uuid = job.ext_runtime_uuid - ext_runtime_name = f'{base_docker_image.split(":")[0]}:{uuid}' - # update job with new extended runtime name - job.runtime_name = ext_runtime_name +def extend_runtime(job, compute_handler, internal_storage): + """ + Points the job to a runtime that bundles its function, building and + deploying it if it does not exist yet. Used when the + runtime_include_function config option is active + """ + base_docker_image = job.runtime_name + job.runtime_name = ( + f'{base_docker_image.split(":")[0]}:{job.ext_runtime_uuid}' + ) - runtime_key = compute_handler.get_runtime_key(job.runtime_name, job.runtime_memory, __version__) + runtime_key = compute_handler.get_runtime_key( + job.runtime_name, job.runtime_memory, __version__ + ) runtime_meta = internal_storage.get_runtime_meta(runtime_key) if not runtime_meta: - ext_docker_file = '/'.join([job.local_tmp_dir, "Dockerfile"]) - - # Generate Dockerfile extended with function dependencies and function - with open(ext_docker_file, 'w') as df: - df.write('\n'.join([ - f'FROM {base_docker_image}', - f'ENV PYTHONPATH={SA_INSTALL_DIR}/modules:$PYTHONPATH', - f'COPY . {SA_INSTALL_DIR}' - ])) - - # Build new extended runtime tagged by function hash - cwd = os.getcwd() - os.chdir(job.local_tmp_dir) - compute_handler.build_runtime(ext_runtime_name, ext_docker_file) - os.chdir(cwd) - shutil.rmtree(job.local_tmp_dir, ignore_errors=True) - - runtime_meta = compute_handler.deploy_runtime(ext_runtime_name, job.runtime_memory, job.runtime_timeout) + _build_extended_runtime(job, compute_handler, base_docker_image) + runtime_meta = compute_handler.deploy_runtime( + job.runtime_name, + job.runtime_memory, + job.runtime_timeout, + ) runtime_meta['runtime_timeout'] = job.runtime_timeout internal_storage.put_runtime_meta(runtime_key, runtime_meta) - # Verify python version and lithops version - if __version__ != runtime_meta['lithops_version']: - raise Exception( - f"Lithops version mismatch. Host version: {__version__} - " - f"Runtime version: {runtime_meta['lithops_version']}" - ) - - py_local_version = version_str(sys.version_info) - py_remote_version = runtime_meta['python_version'] - if py_local_version != py_remote_version: - raise Exception( - f"The runtime '{job.runtime_name}' uses Python {py_remote_version}, " - f"which is incompatible with local Python {py_local_version}" - ) + _verify_runtime_meta(runtime_meta, job.runtime_name) diff --git a/lithops/job/__init__.py b/lithops/job/__init__.py index 3f624091e..6df75691d 100644 --- a/lithops/job/__init__.py +++ b/lithops/job/__init__.py @@ -1,5 +1,4 @@ -from .job import create_map_job -from .job import create_reduce_job +from .job import create_map_job, create_reduce_job __all__ = [ 'create_map_job', diff --git a/lithops/job/job.py b/lithops/job/job.py index eeb6f44b4..da2702161 100644 --- a/lithops/job/job.py +++ b/lithops/job/job.py @@ -22,61 +22,159 @@ import inspect import pickle import logging +import weakref +from collections.abc import Callable, Iterable, Mapping from types import SimpleNamespace +from typing import Any, Dict, List, Optional, Set, Tuple from lithops import utils from lithops.job.partitioner import create_partitions from lithops.storage.utils import create_func_key, create_data_key, \ create_job_key, func_key_suffix -from lithops.job.serialize import SerializeIndependent, create_module_data -from lithops.constants import MAX_AGG_DATA_SIZE, LOCALHOST, \ - SERVERLESS, STANDALONE, CUSTOM_RUNTIME_DIR +from lithops.job.serialize import ( + SerializeIndependent, create_module_data, write_module_data +) +from lithops.constants import ( + MAX_AGG_DATA_SIZE, SERVERLESS, STANDALONE, CUSTOM_RUNTIME_DIR, JOBS_PREFIX +) logger = logging.getLogger(__name__) FUNCTION_CACHE = set() +_FUNC_SERIALIZE_CACHE = weakref.WeakKeyDictionary() MAX_DATA_IN_PAYLOAD = 8 * 1024 # Per invocation. 8KB +def invalidate_function_cache(executor_id: str) -> None: + """ + Drops the cached func keys of an executor, after they were deleted from + the storage backend and have to be uploaded again + """ + prefix = f'{JOBS_PREFIX}/{executor_id}/' + FUNCTION_CACHE.difference_update( + key for key in tuple(FUNCTION_CACHE) if key.startswith(prefix) + ) + + +def _freeze_module_set(mods: Optional[Set[str]]) -> Optional[Tuple[str, ...]]: + if mods is None: + return None + return tuple(sorted(mods)) + + +def _cached_func_serialize( + serializer: Any, + func: Callable, + inc_modules: Optional[Set[str]], + exc_modules: Set[str] +) -> Tuple[bytes, Set[str]]: + """ + Serializes a function and resolves its modules, reusing the result of a + previous job that ran the same function with the same module filters + """ + subkey = ( + _freeze_module_set(inc_modules), + _freeze_module_set(exc_modules), + ) + try: + per_func = _FUNC_SERIALIZE_CACHE[func] + except TypeError: + # Not every callable can be weak referenced, so caching is best effort + per_func = None + except KeyError: + per_func = {} + try: + _FUNC_SERIALIZE_CACHE[func] = per_func + except TypeError: + per_func = None + + if per_func is not None: + cached = per_func.get(subkey) + if cached is not None: + return cached + + func_str = serializer.dumps([func])[0] + func_paths = serializer.module_paths([func], inc_modules, exc_modules) + cached = (func_str, func_paths) + if per_func is not None: + per_func[subkey] = cached + return cached + + +def _serialize_func_and_data( + serializer: Any, + func: Callable, + iterdata: List, + inc_modules: Optional[Set[str]], + exc_modules: Set[str] +) -> Tuple[bytes, List[bytes], Set[str]]: + """ + Serializes the function apart from its data, so that the function can be + cached across jobs. Falls back to serializing everything in one go for + serializers that only expose a call interface + """ + dumps = getattr(serializer, 'dumps', None) + module_paths = getattr(serializer, 'module_paths', None) + if dumps is None or module_paths is None: + ser, paths = serializer( + [func] + list(iterdata), inc_modules, exc_modules + ) + return ser[0], ser[1:], paths + + func_str, func_paths = _cached_func_serialize( + serializer, func, inc_modules, exc_modules + ) + data_strs = dumps(iterdata) + # The data is only inspected for modules when the module manager is on + # and no explicit include list was given + if inc_modules is not None and not inc_modules: + data_paths = module_paths(iterdata, inc_modules, exc_modules) + else: + data_paths = set() + return func_str, data_strs, func_paths | data_paths + + def create_map_job( - config, - internal_storage, - executor_id, - job_id, - map_function, - iterdata, - runtime_meta, - runtime_memory, - extra_env, - include_modules, - exclude_modules, - execution_timeout, - chunksize=None, - extra_args=None, - obj_chunk_size=None, - obj_newline='\n', - obj_chunk_number=None -): + config: Dict[str, Any], + internal_storage: Any, + executor_id: str, + job_id: str, + map_function: Callable, + iterdata: Iterable, + runtime_meta: Mapping[str, Any], + runtime_memory: Optional[int], + extra_env: Optional[Mapping[str, Any]], + include_modules: Optional[Iterable[str]], + exclude_modules: Optional[Iterable[str]], + execution_timeout: Optional[int], + chunksize: Optional[int] = None, + extra_args: Any = None, + obj_chunk_size: Optional[int] = None, + obj_newline: Optional[str] = '\n', + obj_chunk_number: Optional[int] = None +) -> SimpleNamespace: """ - Wrapper to create a map job. It integrates COS logic to process objects. + Creates a map job, splitting the referenced objects into partitions first + when the function processes data from object storage """ host_job_meta = {'host_job_create_tstamp': time.time()} map_iterdata = utils.verify_args(map_function, iterdata, extra_args) - # Object processing functionality - ppo = None + parts_per_object = None if utils.is_object_processing_function(map_function): create_partitions_start = time.time() - # Create partitions according chunk_size or chunk_number - logger.debug('ExecutorID {} | JobID {} - Calling map on partitions ' - 'from object storage flow'.format(executor_id, job_id)) - map_iterdata, ppo = create_partitions( + logger.debug( + f'{utils.log_prefix(executor_id, job_id)} - Calling map on partitions ' + 'from object storage flow' + ) + map_iterdata, parts_per_object = create_partitions( config, internal_storage, map_iterdata, obj_chunk_size, obj_chunk_number, obj_newline ) - host_job_meta['host_job_create_partitions_time'] = round(time.time() - create_partitions_start, 6) - # ######## + host_job_meta['host_job_create_partitions_time'] = round( + time.time() - create_partitions_start, 6 + ) job = _create_job( config=config, @@ -95,49 +193,47 @@ def create_map_job( host_job_meta=host_job_meta ) - if ppo: - job.parts_per_object = ppo + if parts_per_object: + job.parts_per_object = parts_per_object return job def create_reduce_job( - config, - internal_storage, - executor_id, - reduce_job_id, - reduce_function, - map_job, - map_futures, - runtime_meta, - runtime_memory, - obj_reduce_by_key, - extra_env, - include_modules, - exclude_modules, - execution_timeout=None, - extra_args=None -): + config: Dict[str, Any], + internal_storage: Any, + executor_id: str, + reduce_job_id: str, + reduce_function: Callable, + map_job: Any, + map_futures: List, + runtime_meta: Mapping[str, Any], + runtime_memory: Optional[int], + obj_reduce_by_key: Any, + extra_env: Optional[Mapping[str, Any]], + include_modules: Optional[Iterable[str]], + exclude_modules: Optional[Iterable[str]], + execution_timeout: Optional[int] = None, + extra_args: Any = None +) -> SimpleNamespace: """ - Wrapper to create a reduce job. Apply a function across all map futures. + Creates a reduce job that applies a function over the futures of a map + job, either over all of them at once or over one group per source object """ host_job_meta = {'host_job_create_tstamp': time.time()} iterdata = [(map_futures, )] if hasattr(map_job, 'parts_per_object') and obj_reduce_by_key: - prev_total_partitons = 0 + offset = 0 iterdata = [] for total_partitions in map_job.parts_per_object: - iterdata.append((map_futures[prev_total_partitons:prev_total_partitons + total_partitions],)) - prev_total_partitons += total_partitions + end = offset + total_partitions + iterdata.append((map_futures[offset:end],)) + offset = end - reduce_job_env = {'__LITHOPS_REDUCE_JOB': True} - if extra_env is None: - ext_env = reduce_job_env - else: - ext_env = extra_env.copy() - ext_env.update(reduce_job_env) + ext_env = {} if extra_env is None else extra_env.copy() + ext_env['__LITHOPS_REDUCE_JOB'] = True iterdata = utils.verify_args(reduce_function, iterdata, extra_args) @@ -158,68 +254,31 @@ def create_reduce_job( ) -def _create_job( - config, - internal_storage, - executor_id, - job_id, - func, - iterdata, - runtime_meta, - runtime_memory, - extra_env, - include_modules, - exclude_modules, - execution_timeout, - host_job_meta, - chunksize=None -): - """ - Creates a new Job - """ - ext_env = {} if extra_env is None else extra_env.copy() - if ext_env: - ext_env = utils.convert_bools_to_string(ext_env) - logger.debug("Extra environment vars {}".format(ext_env)) - - mode = config['lithops']['mode'] - backend = config['lithops']['backend'] - - job = SimpleNamespace() - job.chunksize = chunksize or config['lithops']['chunksize'] - job.worker_processes = config[backend]['worker_processes'] - job.execution_timeout = execution_timeout or config['lithops']['execution_timeout'] - job.executor_id = executor_id - job.job_id = job_id - job.job_key = create_job_key(job.executor_id, job.job_id) - job.extra_env = ext_env - job.function_name = func.__name__ if inspect.isfunction(func) or inspect.ismethod(func) else type(func).__name__ - job.total_calls = len(iterdata) - - if mode == SERVERLESS: - job.runtime_memory = runtime_memory or config[backend]['runtime_memory'] - job.runtime_timeout = runtime_meta['runtime_timeout'] - if job.execution_timeout >= job.runtime_timeout: - job.execution_timeout = job.runtime_timeout - 5 - - elif mode in STANDALONE: - job.runtime_memory = None - runtime_timeout = config[STANDALONE]['hard_dismantle_timeout'] - if job.execution_timeout >= runtime_timeout: - job.execution_timeout = runtime_timeout - 10 +def _function_name(func: Callable) -> str: + if inspect.isfunction(func) or inspect.ismethod(func): + return func.__name__ + return type(func).__name__ - elif mode == LOCALHOST: - job.runtime_memory = None - job.runtime_timeout = None +def _include_exclude_modules( + config: Mapping[str, Any], + include_modules: Optional[Iterable[str]], + exclude_modules: Optional[Iterable[str]], +) -> Tuple[Optional[Set[str]], Set[str]]: + """ + Merges the module filters given to the job with the ones in the config. + An include set of None means that no module analysis is done at all + """ exclude_modules_cfg = config['lithops'].get('exclude_modules', []) include_modules_cfg = config['lithops'].get('include_modules', []) - if type(include_modules_cfg) is str: + if isinstance(include_modules_cfg, str): if include_modules_cfg.lower() == 'none': include_modules_cfg = None else: - raise ValueError("'include_modules' parameter in config must be a list") + raise ValueError( + "'include_modules' parameter in config must be a list" + ) exc_modules = set() inc_modules = set() @@ -237,132 +296,280 @@ def _create_job( if include_modules is None: inc_modules = None - logger.debug(f'ExecutorID {executor_id} | JobID {job_id} - Serializing function and data') - job_serialize_start = time.time() + return inc_modules, exc_modules + + +def _apply_mode_limits( + job: SimpleNamespace, + config: Mapping[str, Any], + runtime_meta: Mapping[str, Any], + runtime_memory: Optional[int] +) -> None: + """ + Sets the memory and the timeout the job is allowed, clamping the execution + timeout so that the job ends before the runtime is torn down under it + """ + mode = config['lithops']['mode'] + backend = config['lithops']['backend'] + + if mode == SERVERLESS: + job.runtime_memory = ( + config[backend]['runtime_memory'] + if runtime_memory is None + else runtime_memory + ) + job.runtime_timeout = runtime_meta['runtime_timeout'] + if job.execution_timeout >= job.runtime_timeout: + job.execution_timeout = job.runtime_timeout - 5 + return + + job.runtime_memory = None + if mode == STANDALONE: + runtime_timeout = config[STANDALONE]['hard_dismantle_timeout'] + if job.execution_timeout >= runtime_timeout: + job.execution_timeout = runtime_timeout - 10 + return + + job.runtime_timeout = None + + +def _serialize_job( + runtime_meta: Mapping[str, Any], + func: Callable, + iterdata: List, + inc_modules: Optional[Set[str]], + exc_modules: Set[str], + host_job_meta: Dict[str, Any] +) -> SimpleNamespace: + """ + Serializes the function, its module dependencies and the data, and records + how long it took and how big the result is + """ + serialize_start = time.time() serializer = SerializeIndependent(runtime_meta['preinstalls']) - func_and_data_ser, mod_paths = serializer([func] + iterdata, inc_modules, exc_modules) - data_strs = func_and_data_ser[1:] - data_size_bytes = sum(len(x) for x in data_strs) + func_str, data_strs, mod_paths = _serialize_func_and_data( + serializer, func, iterdata, inc_modules, exc_modules + ) module_data = create_module_data(mod_paths) - func_str = func_and_data_ser[0] - func_module_str = pickle.dumps({'func': func_str, 'module_data': module_data}, -1) - func_module_size_bytes = len(func_module_str) + func_module_str = pickle.dumps( + {'func': func_str, 'module_data': module_data}, -1 + ) + data_size_bytes = sum(len(data_str) for data_str in data_strs) - host_job_meta['host_job_serialize_time'] = round(time.time() - job_serialize_start, 6) + host_job_meta['host_job_serialize_time'] = round( + time.time() - serialize_start, 6 + ) host_job_meta['func_data_size_bytes'] = data_size_bytes - host_job_meta['func_module_size_bytes'] = func_module_size_bytes + host_job_meta['func_module_size_bytes'] = len(func_module_str) + + return SimpleNamespace( + func_str=func_str, + data_strs=data_strs, + data_size_bytes=data_size_bytes, + mod_paths=mod_paths, + module_data=module_data, + func_module_str=func_module_str + ) - # Check data limit - if 'data_limit' in config['lithops']: - data_limit = config['lithops']['data_limit'] - else: - data_limit = MAX_AGG_DATA_SIZE - if data_limit and data_size_bytes > data_limit * 1024**2: - log_msg = ('ExecutorID {} | JobID {} - Total data exceeded maximum size ' - 'of {}'.format(executor_id, job_id, utils.sizeof_fmt(data_limit * 1024**2))) - raise Exception(log_msg) - - # Upload function and data - upload_function = not config[backend].get("runtime_include_function", False) - upload_data = any([(len(data_str) * job.chunksize) > MAX_DATA_IN_PAYLOAD for data_str in data_strs]) - - # Upload function and modules - if upload_function: - function_hash = hashlib.md5(func_module_str).hexdigest() - job.func_key = create_func_key(executor_id, function_hash) - if job.func_key not in FUNCTION_CACHE: - logger.debug('ExecutorID {} | JobID {} - Uploading function and modules ' - 'to the storage backend'.format(executor_id, job_id)) - func_upload_start = time.time() - internal_storage.put_func(job.func_key, func_module_str) - func_upload_end = time.time() - host_job_meta['host_func_upload_time'] = round(func_upload_end - func_upload_start, 6) - FUNCTION_CACHE.add(job.func_key) - else: - logger.debug('ExecutorID {} | JobID {} - Function and modules ' - 'found in local cache'.format(executor_id, job_id)) - host_job_meta['host_func_upload_time'] = 0 - else: - # Prepare function and modules locally to store in the runtime image later - function_file = func.__code__.co_filename - function_hash = hashlib.md5(open(function_file, 'rb').read()).hexdigest()[:16] - mod_hash = hashlib.md5(repr(sorted(mod_paths)).encode('utf-8')).hexdigest()[:16] - job.func_key = func_key_suffix - job.ext_runtime_uuid = f'{function_hash}{mod_hash}' - job.local_tmp_dir = os.path.join(CUSTOM_RUNTIME_DIR, job.ext_runtime_uuid) - _store_func_and_modules(job.local_tmp_dir, job.func_key, func_str, module_data) +def _upload_function( + internal_storage: Any, + job: SimpleNamespace, + serialized: SimpleNamespace, + host_job_meta: Dict[str, Any], + prefix: str +) -> None: + """ + Uploads the serialized function and its modules to the storage backend, + unless this executor already uploaded an identical one + """ + function_hash = hashlib.md5(serialized.func_module_str).hexdigest() + job.func_key = create_func_key(job.executor_id, function_hash) + + if job.func_key in FUNCTION_CACHE: + logger.debug(f'{prefix} - Function and modules found in local cache') host_job_meta['host_func_upload_time'] = 0 + return - # upload data - if upload_data or config['lithops']['backend_type'] == utils.BackendType.BATCH.value: - # Upload iterdata to COS only if a single element is greater than MAX_DATA_IN_PAYLOAD - logger.debug('ExecutorID {} | JobID {} - Uploading data to the storage backend' - .format(executor_id, job_id)) - # pass_iteradata through an object storage file - data_key = create_data_key(executor_id, job_id) - job.data_key = data_key - data_bytes, data_byte_ranges = utils.agg_data(data_strs) - job.data_byte_ranges = data_byte_ranges - data_upload_start = time.time() - internal_storage.put_data(data_key, data_bytes) - data_upload_end = time.time() - host_job_meta['host_data_upload_time'] = round(data_upload_end - data_upload_start, 6) + logger.debug( + f'{prefix} - Uploading function and modules to the storage backend' + ) + upload_start = time.time() + internal_storage.put_func(job.func_key, serialized.func_module_str) + host_job_meta['host_func_upload_time'] = round( + time.time() - upload_start, 6 + ) + FUNCTION_CACHE.add(job.func_key) - else: - # pass iteradata as part of the invocation payload - logger.debug('ExecutorID {} | JobID {} - Data per activation is < ' - '{}. Passing data through invocation payload' - .format(executor_id, job_id, utils.sizeof_fmt(MAX_DATA_IN_PAYLOAD))) + +def _bundle_function_in_runtime( + job: SimpleNamespace, + func: Callable, + serialized: SimpleNamespace, + host_job_meta: Dict[str, Any] +) -> None: + """ + Writes the function and its modules to a local directory, for backends + that build them into the runtime image instead of uploading them + """ + with open(func.__code__.co_filename, 'rb') as fid: + function_hash = hashlib.md5(fid.read()).hexdigest()[:16] + mod_hash = hashlib.md5( + repr(sorted(serialized.mod_paths)).encode('utf-8') + ).hexdigest()[:16] + + job.func_key = func_key_suffix + # The uuid identifies the runtime image that carries this exact function + job.ext_runtime_uuid = f'{function_hash}{mod_hash}' + job.local_tmp_dir = os.path.join(CUSTOM_RUNTIME_DIR, job.ext_runtime_uuid) + _store_func_and_modules( + job.local_tmp_dir, job.func_key, serialized.func_str, + serialized.module_data + ) + host_job_meta['host_func_upload_time'] = 0 + + +def _attach_data( + config: Mapping[str, Any], + internal_storage: Any, + job: SimpleNamespace, + serialized: SimpleNamespace, + host_job_meta: Dict[str, Any], + prefix: str +) -> None: + """ + Uploads the data of the job to the storage backend, or leaves it in the + job so that it travels inside the invocation payload when it is small + """ + fits_in_payload = all( + (len(data_str) * job.chunksize) <= MAX_DATA_IN_PAYLOAD + for data_str in serialized.data_strs + ) + is_batch = ( + config['lithops']['backend_type'] == utils.BackendType.BATCH.value + ) + + if fits_in_payload and not is_batch: + logger.debug( + f'{prefix} - Data per activation is < ' + f'{utils.sizeof_fmt(MAX_DATA_IN_PAYLOAD)}. ' + 'Passing data through invocation payload' + ) job.data_key = None job.data_byte_ranges = None - job.data_byte_strs = data_strs + job.data_byte_strs = serialized.data_strs host_job_meta['host_data_upload_time'] = 0 + return + + logger.debug(f'{prefix} - Uploading data to the storage backend') + job.data_key = create_data_key(job.executor_id, job.job_id) + data_bytes, job.data_byte_ranges = utils.agg_data(serialized.data_strs) + upload_start = time.time() + internal_storage.put_data(job.data_key, data_bytes) + host_job_meta['host_data_upload_time'] = round( + time.time() - upload_start, 6 + ) - host_job_meta['host_job_created_time'] = round(time.time() - host_job_meta['host_job_create_tstamp'], 6) +def _create_job( + config: Dict[str, Any], + internal_storage: Any, + executor_id: str, + job_id: str, + func: Callable, + iterdata: List, + runtime_meta: Mapping[str, Any], + runtime_memory: Optional[int], + extra_env: Optional[Mapping[str, Any]], + include_modules: Optional[Iterable[str]], + exclude_modules: Optional[Iterable[str]], + execution_timeout: Optional[int], + host_job_meta: Dict[str, Any], + chunksize: Optional[int] = None +) -> SimpleNamespace: + """ + Creates a new job, uploading its function and its data so that the + invoker only has to hand the workers a reference to them + """ + ext_env = {} if extra_env is None else extra_env.copy() + if ext_env: + ext_env = utils.convert_bools_to_string(ext_env) + logger.debug(f'Extra environment vars {ext_env}') + + backend = config['lithops']['backend'] + prefix = utils.log_prefix(executor_id, job_id) + + job = SimpleNamespace() + job.chunksize = ( + config['lithops']['chunksize'] if chunksize is None else chunksize + ) + job.worker_processes = config[backend]['worker_processes'] + job.execution_timeout = ( + config['lithops']['execution_timeout'] + if execution_timeout is None + else execution_timeout + ) + job.executor_id = executor_id + job.job_id = job_id + job.job_key = create_job_key(job.executor_id, job.job_id) + job.extra_env = ext_env + job.function_name = _function_name(func) + job.total_calls = len(iterdata) + + _apply_mode_limits(job, config, runtime_meta, runtime_memory) + + inc_modules, exc_modules = _include_exclude_modules( + config, include_modules, exclude_modules + ) + + logger.debug(f'{prefix} - Serializing function and data') + serialized = _serialize_job( + runtime_meta, func, iterdata, inc_modules, exc_modules, host_job_meta + ) + + data_limit = config['lithops'].get('data_limit', MAX_AGG_DATA_SIZE) + if data_limit and serialized.data_size_bytes > data_limit * 1024**2: + raise Exception( + f'{prefix} - Total data exceeded maximum size ' + f'of {utils.sizeof_fmt(data_limit * 1024**2)}' + ) + + if config[backend].get('runtime_include_function', False): + _bundle_function_in_runtime(job, func, serialized, host_job_meta) + else: + _upload_function( + internal_storage, job, serialized, host_job_meta, prefix + ) + + _attach_data( + config, internal_storage, job, serialized, host_job_meta, prefix + ) + + host_job_meta['host_job_created_time'] = round( + time.time() - host_job_meta['host_job_create_tstamp'], 6 + ) job.metadata = host_job_meta return job def _store_func_and_modules( - job_tmp_dir, - func_key, - func_str, - module_data -): - ''' stores function and modules in temporary directory to be - used later in optimized runtime - ''' - # save function + job_tmp_dir: str, + func_key: str, + func_str: bytes, + module_data: Optional[Dict[str, str]] +) -> None: + """ + Stores a function and its modules in a local directory, for the custom + runtime build to pick them up + """ os.makedirs(job_tmp_dir, exist_ok=True) - with open(os.path.join(job_tmp_dir, func_key), "wb") as f: - pickle.dump({'func': func_str}, f, -1) + with open(os.path.join(job_tmp_dir, func_key), 'wb') as fid: + pickle.dump({'func': func_str}, fid, -1) - # save modules if module_data: - logger.debug("Writing Function dependencies to local disk") - - modules_path = '/'.join([job_tmp_dir, 'modules']) - - for m_filename, m_data in module_data.items(): - m_path = os.path.dirname(m_filename) - - if len(m_path) > 0 and m_path[0] == "/": - m_path = m_path[1:] - to_make = os.path.join(modules_path, m_path) - try: - os.makedirs(to_make) - except OSError as e: - if e.errno == 17: - pass - else: - raise e - full_filename = os.path.join(to_make, os.path.basename(m_filename)) - - with open(full_filename, 'wb') as fid: - fid.write(utils.b64str_to_bytes(m_data)) + logger.debug('Writing Function dependencies to local disk') + write_module_data(os.path.join(job_tmp_dir, 'modules'), module_data) - logger.debug("Finished storing function and modules") + logger.debug('Finished storing function and modules') diff --git a/lithops/job/partitioner.py b/lithops/job/partitioner.py index f2ec4ed2e..3b8f37d50 100644 --- a/lithops/job/partitioner.py +++ b/lithops/job/partitioner.py @@ -16,9 +16,11 @@ # import os +import posixpath import logging import requests from concurrent.futures import ThreadPoolExecutor +from typing import Any, Callable, Dict, List, Optional, Tuple from lithops import utils from lithops.storage import Storage @@ -29,81 +31,175 @@ CHUNK_THRESHOLD = 128 * 1024 # 128KB +# Backends whose listing supports a glob pattern +_GLOBBER_BACKENDS = ('aws_s3', 'ibm_cos') + +# One entry of the map iterdata, and the partitions it was split into +Entry = Dict[str, Any] +Partitions = Tuple[List[Entry], int] + def create_partitions( - config, - internal_storage, - map_iterdata, - obj_chunk_size, - obj_chunk_number, - obj_newline -): + config: Dict[str, Any], + internal_storage: Any, + map_iterdata: List[Entry], + obj_chunk_size: Optional[int], + obj_chunk_number: Optional[int], + obj_newline: Optional[str] +) -> Tuple[List[Entry], List[int]]: """ - Method that returns the function that will create - the partitions of the objects in the Cloud + Splits the objects referenced by the iterdata into partitions, one task + each. Only one kind of source is partitioned per call """ - urls = [] paths = [] objects = [] logger.debug("Parsing input data") - # first filter; decide if the iterdata elements are urls, - # paths or object storage objects for elem in map_iterdata: if str(elem['obj']).startswith('http'): - # iterdata is a list of public urls urls.append(elem) - elif str(elem['obj']).startswith('/'): - # iterdata is a list of localhost paths or dirs paths.append(elem) - else: - # assume iterdata contains buckets or object keys objects.append(elem) if urls: - # process objects from urls. return _split_objects_from_urls( - urls, obj_chunk_size, - obj_chunk_number, obj_newline + urls, obj_chunk_size, obj_chunk_number, obj_newline ) - - elif paths: - # process objects from localhost paths. + if paths: return _split_objects_from_paths( - paths, obj_chunk_size, - obj_chunk_number, obj_newline + paths, obj_chunk_size, obj_chunk_number, obj_newline ) - - elif objects: - # process objects from an object store. + if objects: return _split_objects_from_object_storage( objects, obj_chunk_size, obj_chunk_number, internal_storage, config, obj_newline ) + return [], [] -def _split_objects_from_urls( - map_func_args_list, - chunk_size, - chunk_number, - obj_newline -): - """ - Create partitions from a list of objects urls - """ +def _log_chunk_settings(chunk_size: Optional[int], chunk_number: Optional[int]) -> None: if chunk_number: - logger.debug(f'Chunk size set to {chunk_size}') - elif chunk_size: logger.debug(f'Chunk number set to {chunk_number}') + elif chunk_size: + logger.debug(f'Chunk size set to {chunk_size}') else: - logger.debug('Chunk size and chunk number not set ') + logger.debug('Chunk size and chunk number not set') + + +def _chunk_from_number(obj_size: int, chunk_number: int) -> int: + chunk_rest = obj_size % chunk_number + return (obj_size // chunk_number) + round((chunk_rest / chunk_number) + 0.5) + + +def _sized_object_chunk( + obj_size: Optional[int], + chunk_number: Optional[int], + chunk_size: Optional[int] +) -> Optional[int]: + """ + Resolves the chunk size to use for one object. A requested chunk number + wins over a requested chunk size, and an unset one means a single chunk + """ + if chunk_number and obj_size: + return _chunk_from_number(obj_size, chunk_number) + if chunk_size and obj_size: + return chunk_size + if obj_size: + return obj_size + return None + + +def _build_partitions( + entry: Entry, + make_obj: Callable[[], Any], + obj_size: Optional[int], + obj_chunk_size: Optional[int], + obj_newline: Optional[str], + label: str +) -> Partitions: + """ + Splits one object into as many partitions as its chunk size calls for, + each one a copy of the entry carrying its own byte range + """ + if not obj_size or not obj_chunk_size: + return [], 0 + + obj_partitions = [] + size = obj_total_partitions = 0 + + parts = obj_size // obj_chunk_size + (obj_size % obj_chunk_size > 0) + logger.debug( + f'Creating {parts} partitions from {label} ({sizeof_fmt(obj_size)})' + ) + + while size < obj_size: + if obj_size <= obj_chunk_size: + # A single partition reads the whole object, no range needed + brange = None + obj_chunk_size = obj_size + elif obj_newline is None: + brange = (size, size + obj_chunk_size - 1) + elif size + obj_chunk_size < obj_size: + # Records must not be cut in two: a partition starts one byte + # early to see whether it begins mid record, and overshoots by + # CHUNK_THRESHOLD so that it can finish the last record it reads + brange = ( + size - 1 if size > 0 else 0, + size + obj_chunk_size + CHUNK_THRESHOLD + ) + else: + brange = (size - 1, obj_size - 1) + obj_chunk_size = obj_size - size + + obj_total_partitions += 1 + + partition = entry.copy() + partition['obj'] = make_obj() + partition['obj'].data_byte_range = brange + partition['obj'].chunk_size = obj_chunk_size + partition['obj'].part = obj_total_partitions + partition['obj'].newline = obj_newline + obj_partitions.append(partition) + + size += obj_chunk_size + + # Only known once the loop is over, so it is filled in afterwards + for partition in obj_partitions: + partition['obj'].total_parts = obj_total_partitions + + return obj_partitions, obj_total_partitions + +def _collect_partitions( + split_fn: Callable[[Entry], Partitions], + entries: List[Entry] +) -> Tuple[List[Entry], List[int]]: + """ + Splits every entry in parallel, since sizing an object needs a request + """ partitions = [] parts_per_object = [] + with ThreadPoolExecutor(64) as ex: + for obj_partitions, nparts in ex.map(split_fn, entries): + partitions.extend(obj_partitions) + parts_per_object.append(nparts) + return partitions, parts_per_object + + +def _split_objects_from_urls( + map_func_args_list: List[Entry], + chunk_size: Optional[int], + chunk_number: Optional[int], + obj_newline: Optional[str] +) -> Tuple[List[Entry], List[int]]: + """ + Creates partitions from a list of object URLs + """ + _log_chunk_settings(chunk_size, chunk_number) def _split(entry): obj_size = None @@ -113,326 +209,232 @@ def _split(entry): if 'content-length' in metadata.headers: obj_size = int(metadata.headers['content-length']) - if chunk_number and obj_size: - chunk_rest = obj_size % chunk_number - obj_chunk_size = (obj_size // chunk_number) + \ - round((chunk_rest / chunk_number) + 0.5) - elif chunk_size and obj_size: - obj_chunk_size = chunk_size - elif obj_size: - obj_chunk_size = obj_size - else: + obj_chunk_size = _sized_object_chunk(obj_size, chunk_number, chunk_size) + if obj_size is None: + # Size unknown, so the object is a single partition of unset size obj_chunk_size = obj_size = 1 if 'accept-ranges' not in metadata.headers: obj_chunk_size = obj_size - obj_partitions = [] - size = obj_total_partitions = 0 - - ci = obj_size - cz = obj_chunk_size - parts = ci // cz + (ci % cz > 0) - logger.debug(f'Creating {parts} partitions from url {object_url} ({sizeof_fmt(obj_size)})') - - while size < obj_size - 1: - if obj_size <= obj_chunk_size: - # Only one chunk - brange = None - obj_chunk_size = obj_size - elif obj_newline is None: - # partitions of the same size - brange = (size, size + obj_chunk_size - 1) - elif size + obj_chunk_size < obj_size: - # common chunk - brange = (size - 1 if size > 0 else 0, size + obj_chunk_size + CHUNK_THRESHOLD) - else: - # last chunk - brange = (size - 1, obj_size - 1) - obj_chunk_size = obj_size - size - - obj_total_partitions += 1 - - partition = entry.copy() - partition['obj'] = CloudObjectUrl(object_url) - partition['obj'].data_byte_range = brange - partition['obj'].chunk_size = obj_chunk_size - partition['obj'].part = obj_total_partitions - partition['obj'].newline = obj_newline - obj_partitions.append(partition) - - size += obj_chunk_size - - for partition in obj_partitions: - partition['obj'].total_parts = obj_total_partitions - - partitions.extend(obj_partitions) - parts_per_object.append(obj_total_partitions) - - with ThreadPoolExecutor(64) as ex: - ex.map(_split, map_func_args_list) + return _build_partitions( + entry, + lambda: CloudObjectUrl(object_url), + obj_size, + obj_chunk_size, + obj_newline, + f'url {object_url}' + ) - return partitions, parts_per_object + return _collect_partitions(_split, map_func_args_list) -def _split_objects_from_paths( - map_func_args_list, - chunk_size, - chunk_number, - obj_newline -): +def _expand_paths(map_func_args_list: List[Entry]) -> List[Entry]: """ - Create partitions from a list of objects paths + Replaces every directory entry with one entry per file it contains, + dropping duplicates and anything that is not a file """ - if chunk_number: - logger.debug(f'Chunk size set to {chunk_size}') - elif chunk_size: - logger.debug(f'Chunk number set to {chunk_number}') - else: - logger.debug('Chunk size and chunk number not set ') - - partitions = [] - parts_per_object = [] - files = set() - new_map_func_args_list = [] + expanded = [] for elem in map_func_args_list: if os.path.isdir(elem['obj']): path = elem['obj'] - found_files = os.listdir(path) - for filename in found_files: + for filename in os.listdir(path): full_path = os.path.join(path, filename) - if full_path in files or \ - not os.path.isfile(full_path): + if full_path in files or not os.path.isfile(full_path): continue files.add(full_path) new_elem = elem.copy() new_elem['obj'] = full_path - new_map_func_args_list.append(new_elem) + expanded.append(new_elem) elif os.path.isfile(elem['obj']): if elem['obj'] in files: continue files.add(elem['obj']) - new_map_func_args_list.append(elem) + expanded.append(elem) + + return expanded + + +def _split_objects_from_paths( + map_func_args_list: List[Entry], + chunk_size: Optional[int], + chunk_number: Optional[int], + obj_newline: Optional[str] +) -> Tuple[List[Entry], List[int]]: + """ + Creates partitions from a list of local files and directories + """ + _log_chunk_settings(chunk_size, chunk_number) def _split(entry): path = entry['obj'] - file_stats = os.stat(entry['obj']) - obj_size = int(file_stats.st_size) - - if chunk_number and obj_size: - chunk_rest = obj_size % chunk_number - obj_chunk_size = (obj_size // chunk_number) + \ - round((chunk_rest / chunk_number) + 0.5) - elif chunk_size and obj_size: - obj_chunk_size = chunk_size - elif obj_size: - obj_chunk_size = obj_size - else: - obj_chunk_size = obj_size = 1 + obj_size = int(os.stat(path).st_size) + obj_chunk_size = _sized_object_chunk(obj_size, chunk_number, chunk_size) + return _build_partitions( + entry, + lambda: CloudObjectLocal(path), + obj_size, + obj_chunk_size, + obj_newline, + f'path {path}' + ) - obj_partitions = [] - size = obj_total_partitions = 0 - - ci = obj_size - cz = obj_chunk_size - parts = ci // cz + (ci % cz > 0) - logger.debug(f'Creating {parts} partitions from url {path} ({sizeof_fmt(obj_size)})') - - while size < obj_size - 1: - if obj_size <= obj_chunk_size: - # Only one chunk - brange = None - obj_chunk_size = obj_size - elif obj_newline is None: - # partitions of the same size - brange = (size, size + obj_chunk_size - 1) - elif size + obj_chunk_size < obj_size: - # common chunk - brange = (size - 1 if size > 0 else 0, size + obj_chunk_size + CHUNK_THRESHOLD) - else: - # last chunk - brange = (size - 1, obj_size - 1) - obj_chunk_size = obj_size - size - - obj_total_partitions += 1 - - partition = entry.copy() - partition['obj'] = CloudObjectLocal(path) - partition['obj'].data_byte_range = brange - partition['obj'].chunk_size = obj_chunk_size - partition['obj'].part = obj_total_partitions - partition['obj'].newline = obj_newline - obj_partitions.append(partition) - - size += obj_chunk_size - - for partition in obj_partitions: - partition['obj'].total_parts = obj_total_partitions - - partitions.extend(obj_partitions) - parts_per_object.append(obj_total_partitions) + return _collect_partitions(_split, _expand_paths(map_func_args_list)) - with ThreadPoolExecutor(64) as ex: - ex.map(_split, new_map_func_args_list) - return partitions, parts_per_object +def _glob_list_prefix(prefix: str, obj_name: str) -> str: + """Return the listing prefix truncated at the first glob character.""" + if '*' in prefix: + return prefix[:prefix.index('*')] + glob_tail = obj_name[:obj_name.index('*')] + return f'{prefix}/{glob_tail}' if prefix else glob_tail -def _split_objects_from_object_storage( - map_func_args_list, - chunk_size, - chunk_number, - internal_storage, - config, - obj_newline -): +def _resolve_object_storage( + map_func_args_list: List[Entry], + internal_storage: Any, + config: Dict[str, Any] +) -> Any: """ - Create partitions from a list of buckets or object keys + Rewrites every entry as a full storage URL, and returns the client they + all point to. Only one storage backend is supported per map call """ - if chunk_number: - logger.debug(f'Chunk size set to {chunk_size}') - elif chunk_size: - logger.debug(f'Chunk number set to {chunk_number}') - else: - logger.debug('Chunk size and chunk number not set') + backends = set() - sbs = set() - - # check that only one schemma provided. Throw exception if more than one provided for elem in map_func_args_list: - if type(elem['obj']) is CloudObject: - elem['obj'] = f"{elem['obj'].backend}://{elem['obj'].bucket}/{elem['obj'].key}" - sb, bucket, prefix, obj_name = utils.split_object_url(elem['obj']) + if isinstance(elem['obj'], CloudObject): + elem['obj'] = ( + f"{elem['obj'].backend}://{elem['obj'].bucket}/{elem['obj'].key}" + ) + sb, _, _, _ = utils.split_object_url(elem['obj']) if sb is None: sb = internal_storage.backend elem['obj'] = f"{sb}://{elem['obj']}" + backends.add(sb) - sbs.add(sb) - - if len(sbs) > 1: - raise Exception('Process objects from multiple storage backends is not supported. ' - f'Current storage backends: {sbs}') + if len(backends) > 1: + raise Exception( + 'Process objects from multiple storage backends is not supported. ' + f'Current storage backends: {backends}' + ) - sb = sbs.pop() + sb = backends.pop() if sb == internal_storage.backend: - storage = internal_storage.storage - else: - storage = Storage(config=config, backend=sb) - partitions = [] - parts_per_object = [] + return internal_storage.storage + return Storage(config=config, backend=sb) - def _split(bucket, key, entry, obj_size): - if key.endswith('/'): - logger.debug(f'Discarding object "{key}" as it is a prefix folder (0.0B)') - return - - if chunk_number: - chunk_rest = obj_size % chunk_number - obj_chunk_size = (obj_size // chunk_number) + \ - round((chunk_rest / chunk_number) + 0.5) - elif chunk_size: - obj_chunk_size = chunk_size + +def _list_objects( + storage: Any, + sb: str, + bucket: str, + prefix: str, + obj_name: str +) -> List[Dict[str, Any]]: + """ + Lists the metadata of the objects one entry refers to, be it a single + key, a glob pattern, a prefix or a whole bucket + """ + if obj_name: + match_pattern = None + if sb in _GLOBBER_BACKENDS and ('*' in prefix or '*' in obj_name): + match_pattern = posixpath.join(prefix, obj_name) + prefix = _glob_list_prefix(prefix, obj_name) + + prefix = prefix + '/' if prefix else prefix + if match_pattern is not None: + logger.debug( + f"Listing objects with Globber {match_pattern} " + f"in {sb}://{'/'.join([bucket, prefix])}" + ) + return storage.list_objects(bucket, prefix, match_pattern) + + logger.debug( + f"Head on object {sb}://{'/'.join([bucket, prefix, obj_name])}" + ) + object_key = posixpath.join(prefix, obj_name) + head_md = storage.head_object(bucket, object_key) + content_length = head_md.get('content-length') + if content_length is None: + raise KeyError( + f"The {sb} backend reported no content-length for " + f"{object_key}, so its size is unknown" + ) + head_md['Key'] = object_key + head_md['Size'] = int(content_length) + return [head_md] + + if prefix: + match_pattern = None + if sb in _GLOBBER_BACKENDS and '*' in prefix: + match_pattern = prefix + prefix = prefix[:prefix.index('*')] + logger.debug( + f"Listing prefixes with Globber {match_pattern} " + f"in {sb}://{'/'.join([bucket, prefix])}" + ) else: - obj_chunk_size = obj_size + logger.debug( + f"Listing prefixes in {sb}://{'/'.join([bucket, prefix])}" + ) + + prefix = prefix + '/' if prefix else prefix + return storage.list_objects(bucket, prefix, match_pattern) + + logger.debug(f"Listing objects in {sb}://{bucket}") + return storage.list_objects(bucket) + + +def _split_objects_from_object_storage( + map_func_args_list: List[Entry], + chunk_size: Optional[int], + chunk_number: Optional[int], + internal_storage: Any, + config: Dict[str, Any], + obj_newline: Optional[str] +) -> Tuple[List[Entry], List[int]]: + """ + Creates partitions from a list of buckets, prefixes or object keys + """ + _log_chunk_settings(chunk_size, chunk_number) + storage = _resolve_object_storage( + map_func_args_list, internal_storage, config + ) + + partitions = [] + parts_per_object = [] + total_objects = 0 - obj_partitions = [] - size = obj_total_partitions = 0 - - ci = obj_size - cz = obj_chunk_size - parts = ci // cz + (ci % cz > 0) - logger.debug(f'Creating {parts} partitions from object {key} ({sizeof_fmt(obj_size)})') - - while size < obj_size - 1: - if obj_size <= obj_chunk_size: - # Only one chunk - brange = None - obj_chunk_size = obj_size - elif obj_newline is None: - # partitions of the same size - brange = (size, size + obj_chunk_size - 1) - elif size + obj_chunk_size < obj_size: - # common chunk - brange = (size - 1 if size > 0 else 0, size + obj_chunk_size + CHUNK_THRESHOLD) - else: - # last chunk - brange = (size - 1, obj_size - 1) - obj_chunk_size = obj_size - size - - obj_total_partitions += 1 - - partition = entry.copy() - partition['obj'] = CloudObject(sb, bucket, key) - partition['obj'].data_byte_range = brange - partition['obj'].chunk_size = obj_chunk_size - partition['obj'].part = obj_total_partitions - partition['obj'].newline = obj_newline - obj_partitions.append(partition) - - size += obj_chunk_size - - for partition in obj_partitions: - partition['obj'].total_parts = obj_total_partitions - - partitions.extend(obj_partitions) - parts_per_object.append(obj_total_partitions) - - total_objects = int(0) for elem in map_func_args_list: - objects = [] - exclude = {'obj'} - params = {k: elem[k] for k in set(list(elem.keys())) - set(exclude)} + params = {k: v for k, v in elem.items() if k != 'obj'} sb, bucket, prefix, obj_name = utils.split_object_url(elem['obj']) + objects = _list_objects(storage, sb, bucket, prefix, obj_name) - if obj_name: - match_pattern = None - if sb in ['aws_s3', 'ibm_cos'] and (prefix.find('*') > -1 or obj_name.find('*') > -1): - - match_pattern = os.path.join(prefix, obj_name) - - if prefix.find('*') > -1: - prefix = prefix[:prefix.index('*')] - else: - prefix = '/'.join(prefix, obj_name[:obj_name.index('*')]) - - prefix = prefix + '/' if prefix else prefix - if match_pattern is not None: - logger.debug(f"Listing objects with Globber {match_pattern} in {sb}://{'/'.join([bucket, prefix])}") - objects = storage.list_objects(bucket, prefix, match_pattern) - else: - # this is wrong to list prefix only, as it may return more objects than requested - logger.debug(f"Head on object {sb}://{'/'.join([bucket, prefix, obj_name])}") - head_md = storage.head_object(bucket, os.path.join(prefix, obj_name)) - head_md['Key'] = os.path.join(prefix, obj_name) - head_md['Size'] = int(head_md['content-length']) - objects.append(head_md) - - elif prefix: - match_pattern = None - if sb in ['aws_s3', 'ibm_cos'] and prefix.find('*') > -1: - - match_pattern = prefix - if prefix.find('*') > -1: - prefix = prefix[:prefix.index('*')] - - logger.debug(f"Listing prefixes with Globber {match_pattern} in {sb}://{'/'.join([bucket, prefix])}") - else: - logger.debug(f"Listing prefixes in {sb}://{'/'.join([bucket, prefix])}") - - prefix = prefix + '/' if prefix else prefix - objects = storage.list_objects(bucket, prefix, match_pattern) - else: - logger.debug(f"Listing objects in {sb}://{bucket}") - objects = storage.list_objects(bucket) - - total_objects = total_objects + len(objects) for dobj in objects: key = dobj['Key'] - entry = {'obj': f'{sb}://{bucket}/{key}'} - entry.update(params) - _split(bucket, key, entry, dobj['Size']) + if key.endswith('/'): + logger.debug( + f'Discarding object "{key}" as it is a prefix folder (0.0B)' + ) + continue + + total_objects += 1 + obj_size = dobj['Size'] + obj_chunk_size = _sized_object_chunk( + obj_size, chunk_number, chunk_size + ) + obj_partitions, nparts = _build_partitions( + {'obj': f'{sb}://{bucket}/{key}', **params}, + lambda: CloudObject(sb, bucket, key), + obj_size, + obj_chunk_size, + obj_newline, + f'object {key}' + ) + partitions.extend(obj_partitions) + parts_per_object.append(nparts) logger.debug(f"Total objects found: {total_objects}") if total_objects == 0: diff --git a/lithops/job/serialize.py b/lithops/job/serialize.py index 75efc570e..58fbd67cd 100644 --- a/lithops/job/serialize.py +++ b/lithops/job/serialize.py @@ -17,6 +17,7 @@ # import os +import posixpath import glob import importlib import logging @@ -27,188 +28,296 @@ from functools import partial, reduce from importlib import import_module from types import CodeType, FunctionType, ModuleType +from typing import Any, Dict, Iterable, List, Optional, Set, Tuple from lithops.libs import imp from lithops.libs import inspect as linspect -from lithops.utils import bytes_to_b64str +from lithops.utils import bytes_to_b64str, b64str_to_bytes from lithops.libs.multyvac.module_dependency import ModuleDependencyAnalyzer logger = logging.getLogger(__name__) +_BUILTIN_MODULES = {'__builtin__', 'builtins'} -class SerializeIndependent: - def __init__(self, preinstalls): - self.preinstalled_modules = preinstalls - self.preinstalled_modules.append(['lithops', True]) +def _is_user_function(obj: Any) -> bool: + return inspect.isfunction(obj) or ( + inspect.ismethod(obj) and inspect.isfunction(obj.__func__) + ) + + +def _joined_or_none(names: Set[str]) -> Optional[str]: + return ", ".join(names) if names else None + + +def write_module_data(dest_dir: str, module_data: Optional[Dict[str, str]]) -> None: + """ + Writes the encoded module payloads into dest_dir, recreating the package + directories they belong to + """ + if not module_data: + return + for m_filename, m_data in module_data.items(): + # The keys are posix paths built on the client, which may be a + # different platform than the one unpacking them + posix_name = m_filename.replace('\\', '/').lstrip('/') + parent = posixpath.dirname(posix_name) + dest_subdir = ( + os.path.join(dest_dir, *parent.split('/')) if parent else dest_dir + ) + os.makedirs(dest_subdir, exist_ok=True) + full_filename = os.path.join( + dest_subdir, posixpath.basename(posix_name) + ) + with open(full_filename, 'wb') as fid: + fid.write(b64str_to_bytes(m_data)) + + +class SerializeIndependent: + """ + Serializes the function and the data of a job, and finds the modules they + depend on and the runtime does not already provide + """ + + def __init__(self, preinstalls: List): + # Lithops is always in the runtime, even when it is not preinstalled + self.preinstalled_modules = list(preinstalls) + [['lithops', True]] self._modulemgr = None - def __call__(self, list_of_objs, include_modules, exclude_modules): + def dumps(self, list_of_objs: List) -> List[bytes]: + """Serializes every object on its own, so they can be split apart""" + return [cloudpickle.dumps(obj) for obj in list_of_objs] + + def _preinstalled_names(self) -> Set[str]: + return {name for name, _ in self.preinstalled_modules} + + def _referenced_module_paths( + self, + list_of_objs: List, + exclude_modules: Iterable[str] + ) -> Set[str]: """ - Serialize f, args, kwargs independently + Finds the paths of the modules the objects reference, leaving out the + ones the runtime already has and the ones explicitly excluded """ - preinstalled_modules = [name for name, _ in self.preinstalled_modules] - - strs = [] - mod_paths = set() + self._modulemgr = ModuleDependencyAnalyzer() + self._modulemgr.ignore(self._preinstalled_names()) + self._modulemgr.ignore(exclude_modules) + ref_modules = set() for obj in list_of_objs: - strs.append(cloudpickle.dumps(obj)) - - if include_modules is None: - # If include_modules is explicitly set to None, no module is included - logger.debug('Module manager disabled. Modules to transmit: None') - return (strs, mod_paths) + ref_modules.update(self._module_inspect(obj)) - if len(include_modules) == 0: - # If include_modules is not provided (empty list by default), - # inspect the objects looking for referenced modules - self._modulemgr = ModuleDependencyAnalyzer() - self._modulemgr.ignore(preinstalled_modules) - self._modulemgr.ignore(exclude_modules) + logger.debug(f"Referenced Modules: {_joined_or_none(ref_modules)}") - ref_modules = set() - - for obj in list_of_objs: - ref_modules.update(self._module_inspect(obj)) + mod_paths = set() + for module_name in ref_modules: + if module_name in ['__main__', None]: + continue + try: + mod_spec = importlib.util.find_spec(module_name) + except Exception: + mod_spec = None + + origin = mod_spec.origin if mod_spec else module_name + # Native extensions cannot be analysed any further, so they are + # shipped as they are instead of going through the analyzer + if origin and origin.endswith('.so'): + excluded = ( + origin in exclude_modules + or os.path.basename(origin) in exclude_modules + ) + if not excluded: + mod_paths.add(origin) + else: + self._modulemgr.add(module_name) - logger.debug("Referenced Modules: {}".format(None if not - ref_modules else ", ".join(ref_modules))) + return mod_paths | self._modulemgr.get_and_clear_paths() - for module_name in ref_modules: - if module_name in ['__main__', None]: - continue - try: - mod_spec = importlib.util.find_spec(module_name) - except Exception: - mod_spec = None - - origin = mod_spec.origin if mod_spec else module_name - if origin and origin.endswith('.so'): - if origin not in exclude_modules and \ - os.path.basename(origin) not in exclude_modules: - mod_paths.add(origin) - else: - self._modulemgr.add(module_name) + def _explicit_module_paths(self, include_modules: Iterable[str]) -> Set[str]: + """ + Resolves the paths of the modules the user asked for, given either as + a file path or as an importable module name + """ + preinstalled_names = self._preinstalled_names() + mod_paths = set() - tent_mod_paths = self._modulemgr.get_and_clear_paths() - mod_paths = mod_paths.union(tent_mod_paths) + logger.debug(f"Include Modules: {', '.join(include_modules)}") - else: - # If include_modules is provided, include only the provided list - logger.debug("Include Modules: {}".format(", ".join(include_modules))) - for module_name in include_modules: - if module_name.endswith('.so') or module_name.endswith('.py'): - pathname = os.path.abspath(module_name) - if os.path.isfile(pathname): - logger.debug(f"Module '{module_name}' found in {pathname}") - mod_paths.add(pathname) - else: - logger.debug(f"Could not find module '{module_name}', skipping") - continue - module_root = module_name.split('.')[0] - if module_root in preinstalled_modules: - logger.debug(f"Module '{module_name}' is already installed in the runtime, skipping") - continue - try: - fp, pathname, description = imp.find_module(module_root) + for module_name in include_modules: + if module_name.endswith(('.so', '.py')): + pathname = os.path.abspath(module_name) + if os.path.isfile(pathname): logger.debug(f"Module '{module_name}' found in {pathname}") mod_paths.add(pathname) - except ImportError: - logger.debug(f"Could not find module '{module_name}', skipping") + else: + logger.debug( + f"Could not find module '{module_name}', skipping" + ) + continue + + module_root = module_name.split('.')[0] + if module_root in preinstalled_names: + logger.debug( + f"Module '{module_name}' is already installed " + "in the runtime, skipping" + ) + continue - logger.debug("Modules to transmit: {}".format(None if not mod_paths else ", ".join(mod_paths))) + try: + _, pathname, _ = imp.find_module(module_root) + logger.debug(f"Module '{module_name}' found in {pathname}") + mod_paths.add(pathname) + except ImportError: + logger.debug( + f"Could not find module '{module_name}', skipping" + ) - return (strs, mod_paths) + return mod_paths - def _module_inspect(self, obj): + def module_paths( + self, + list_of_objs: List, + include_modules: Optional[Iterable[str]], + exclude_modules: Iterable[str] + ) -> Set[str]: """ - inspect objects for module dependencies + Collects the paths of the modules that have to travel with the job: + either the ones explicitly included, or the ones its code references """ - worklist = [] - seen = set() - mods = set() + if include_modules is None: + logger.debug('Module manager disabled. Modules to transmit: None') + return set() - if inspect.isfunction(obj) or (inspect.ismethod(obj) and inspect.isfunction(obj.__func__)): - # The obj is the user's function - worklist.append(obj) + if include_modules: + mod_paths = self._explicit_module_paths(include_modules) + else: + mod_paths = self._referenced_module_paths( + list_of_objs, exclude_modules + ) - elif type(obj).__name__ == 'cython_function_or_method': - for k, v in linspect.getmembers_static(obj): - if k == '__globals__': - mods.add(v['__file__']) + logger.debug(f"Modules to transmit: {_joined_or_none(mod_paths)}") - elif type(obj) is dict: - # the obj is the user's iterdata + return mod_paths + + def __call__( + self, + list_of_objs: List, + include_modules: Optional[Iterable[str]], + exclude_modules: Iterable[str] + ) -> Tuple[List[bytes], Set[str]]: + """ + Serializes the objects independently and returns them together with + the paths of the modules they depend on + """ + return ( + self.dumps(list_of_objs), + self.module_paths(list_of_objs, include_modules, exclude_modules), + ) + + def _entry_points(self, obj: Any) -> Tuple[List, Set[str]]: + """ + Returns the user functions to inspect for the given job function, plus + the modules that can only be read off the object itself + """ + if _is_user_function(obj): + return [obj], set() + + if type(obj).__name__ == 'cython_function_or_method': + return [], { + value['__file__'] + for name, value in linspect.getmembers_static(obj) + if name == '__globals__' + } + + if isinstance(obj, dict): + worklist = [] for param in obj.values(): - if type(param).__module__ == "__builtin__": - continue - elif inspect.isfunction(param): - # it is a user defined function + if _is_user_function(param): worklist.append(param) - else: - # it is a user defined class - for k, v in linspect.getmembers_static(param): - if inspect.isfunction(v) or (inspect.ismethod(v) and inspect.isfunction(v.__func__)): - worklist.append(v) - elif isinstance(obj, partial): - found_methods = ["__call__"] - worklist.append(obj.func) - else: - # The obj is the user's function but in form of a class - found_methods = [] - for k, v in linspect.getmembers_static(obj): - if inspect.isfunction(v) or (inspect.ismethod(v) and inspect.isfunction(v.__func__)): - found_methods.append(k) - worklist.append(v) - if "__call__" not in found_methods: - raise ValueError( - "The class you passed as the function to " - 'run must contain the "__call__" method' + continue + if getattr(type(param), '__module__', None) in _BUILTIN_MODULES: + continue + worklist.extend( + value for _, value in linspect.getmembers_static(param) + if _is_user_function(value) ) + return worklist, set() - # The worklist is only used for analyzing functions + if isinstance(obj, partial): + return [obj.func], set() + + worklist = [] + found_methods = [] + for name, value in linspect.getmembers_static(obj): + if _is_user_function(value): + found_methods.append(name) + worklist.append(value) + if "__call__" not in found_methods: + raise ValueError( + "The class you passed as the function to " + 'run must contain the "__call__" method' + ) + return worklist, set() + + def _module_inspect(self, obj: Any) -> Set[str]: + """ + Inspects an object for the modules it depends on, following every + function and code object it references in turn + """ + worklist, mods = self._entry_points(obj) + seen = set() + + # Both worklists are appended to while being iterated on purpose: + # that is how the references are followed to the end for fn in worklist: mods.add(fn.__module__) codeworklist = [fn] cvs = inspect.getclosurevars(fn) - modules = list(cvs.nonlocals.items()) - modules.extend(list(cvs.globals.items())) - - for k, v in modules: - if inspect.ismodule(v): - mods.add(v.__name__) - elif inspect.isfunction(v) and id(v) not in seen: - seen.add(id(v)) - mods.add(v.__module__) - worklist.append(v) - elif hasattr(v, "__module__"): - mods.add(v.__module__) + closure_vars = ( + list(cvs.nonlocals.values()) + list(cvs.globals.values()) + ) + + for value in closure_vars: + if inspect.ismodule(value): + mods.add(value.__name__) + elif inspect.isfunction(value) and id(value) not in seen: + seen.add(id(value)) + mods.add(value.__module__) + worklist.append(value) + elif hasattr(value, "__module__"): + mods.add(value.__module__) for block in codeworklist: - for (k, v) in [self._inner_module_inspect(inst) - for inst in Bytecode(block)]: - if k is None: + for kind, value in ( + self._inner_module_inspect(inst) for inst in Bytecode(block) + ): + if kind is None: continue - if k == "modules": - newmods = [mod.__name__ for mod in v if hasattr(mod, "__name__")] - mods.update(set(newmods)) - elif k == "code" and id(v) not in seen: - seen.add(id(v)) - if hasattr(v, "__module__"): - mods.add(v.__module__) - - if inspect.isfunction(v): - worklist.append(v) - elif inspect.iscode(v): - codeworklist.append(v) - - return {mod_name.split(".")[0] for mod_name in mods} - - def _inner_module_inspect(self, inst): + if kind == "modules": + mods.update( + mod.__name__ for mod in value + if hasattr(mod, "__name__") + ) + elif kind == "code" and id(value) not in seen: + seen.add(id(value)) + if hasattr(value, "__module__"): + mods.add(value.__module__) + + if inspect.isfunction(value): + worklist.append(value) + elif inspect.iscode(value): + codeworklist.append(value) + + # Dynamically built functions and code objects can have a + # __module__ of None, which names no module to ship + return {mod_name.split(".")[0] for mod_name in mods if mod_name} + + def _inner_module_inspect(self, inst: Any) -> Tuple[Optional[str], Any]: """ - get interesting modules refernced within an object + Reads the module or the code object that a single bytecode + instruction refers to """ if inst.opname == "IMPORT_NAME": try: @@ -219,33 +328,40 @@ def _inner_module_inspect(self, inst): except Exception: return (None, None) if inst.opname == "LOAD_GLOBAL": - if inst.argval in globals() and type(globals()[inst.argval]) in [CodeType, FunctionType]: - return ("code", globals()[inst.argval]) - if inst.argval in globals() and type(globals()[inst.argval]) == ModuleType: - return ("modules", [globals()[inst.argval]]) - else: - return (None, None) - if "LOAD_" in inst.opname and type(inst.argval) in [CodeType, FunctionType]: + value = globals().get(inst.argval) + if isinstance(value, (CodeType, FunctionType)): + return ("code", value) + if isinstance(value, ModuleType): + return ("modules", [value]) + return (None, None) + if "LOAD_" in inst.opname and isinstance( + inst.argval, (CodeType, FunctionType) + ): return ("code", inst.argval) return (None, None) -def create_module_data(mod_paths): - +def create_module_data(mod_paths: Iterable[str]) -> Dict[str, str]: + """ + Reads the modules at the given paths and encodes them, keyed by the path + they have to be written to relative to their package root + """ module_data = {} - # load mod paths - for m in mod_paths: - if os.path.isdir(m): - files = glob.glob(os.path.join(m, "**/*.py"), recursive=True) - pkg_root = os.path.abspath(os.path.dirname(m)) + + for mod_path in mod_paths: + pkg_root = os.path.abspath(os.path.dirname(mod_path)) + if os.path.isdir(mod_path): + files = glob.glob( + os.path.join(mod_path, "**/*.py"), recursive=True + ) else: - pkg_root = os.path.abspath(os.path.dirname(m)) - files = [m] - for f in files: - f = os.path.abspath(f) - with open(f, 'rb') as file: - mod_str = file.read() - dest_filename = Path(f[len(pkg_root) + 1:]).as_posix() + files = [mod_path] + + for filename in files: + filename = os.path.abspath(filename) + with open(filename, 'rb') as fid: + mod_str = fid.read() + dest_filename = Path(filename[len(pkg_root) + 1:]).as_posix() module_data[dest_filename] = bytes_to_b64str(mod_str) return module_data diff --git a/lithops/localhost/__init__.py b/lithops/localhost/__init__.py index 0ad09e0e6..a1597e76e 100644 --- a/lithops/localhost/__init__.py +++ b/lithops/localhost/__init__.py @@ -1,10 +1,11 @@ from .v1.localhost import LocalhostHandlerV1 from .v2.localhost import LocalhostHandlerV2 -# Set the default localhost handler +# Callers that do not select a version explicitly get v2 LocalhostHandler = LocalhostHandlerV2 __all__ = [ + 'LocalhostHandler', 'LocalhostHandlerV1', - 'LocalhostHandlerV2' + 'LocalhostHandlerV2', ] diff --git a/lithops/localhost/config.py b/lithops/localhost/config.py index 1587d7342..882a47b21 100644 --- a/lithops/localhost/config.py +++ b/lithops/localhost/config.py @@ -13,49 +13,82 @@ # import os +import posixpath import re import sys from enum import Enum +from typing import Any, Dict + +from lithops.version import __version__ DEFAULT_CONFIG_KEYS = { 'runtime': os.path.basename(sys.executable), - 'worker_processes': os.cpu_count(), + 'worker_processes': os.cpu_count() or 1, } LOCALHOST_EXECUTION_TIMEOUT = 3600 +_WINDOWS_PATH = re.compile(r'^[A-Za-z]:\\.*$') +# Interpreters like python, python3, python3.12, python.exe — not docker tags +# such as python:3.12. +_PYTHON_INTERPRETER = re.compile( + r'^python(\d+(\.\d+)*)?(\.exe)?$', + re.IGNORECASE, +) + -class LocvalhostEnvironment(Enum): +class LocalhostEnvironment(Enum): + """Where a localhost job runs: this Python installation or a container""" DEFAULT = "default" CONTAINER = "container" -def get_environment(runtime_name): - - windows_path_pattern = re.compile(r'^[A-Za-z]:\\.*$') - if runtime_name.startswith(('python', '/')) \ - or windows_path_pattern.match(runtime_name) is not None: - environment = LocvalhostEnvironment.DEFAULT - else: - environment = LocvalhostEnvironment.CONTAINER - - return environment - - -def load_config(config_data): - +def get_environment(runtime_name: str) -> LocalhostEnvironment: + """ + Decides the environment a runtime name refers to. An absolute path or an + interpreter name is run locally, anything else is a container image + """ + basename = os.path.basename(runtime_name) + if ( + runtime_name.startswith('/') + or _WINDOWS_PATH.match(runtime_name) is not None + or _PYTHON_INTERPRETER.match(basename) is not None + ): + return LocalhostEnvironment.DEFAULT + return LocalhostEnvironment.CONTAINER + + +def runtime_key(runtime_name: str) -> str: + """ + Builds the key the runtime metadata is cached under. Always POSIX, so that + a Windows client and a Unix one agree on the same key + """ + name = runtime_name.replace('\\', '/').strip('/') + return posixpath.join('localhost', __version__, name) + + +def runtime_info(config: Dict[str, Any]) -> Dict[str, Any]: + """Returns the runtime limits the executor reports to the user""" + return { + 'runtime_name': config['runtime'], + 'runtime_memory': config.get('runtime_memory'), + 'runtime_timeout': config.get('runtime_timeout'), + 'max_workers': config['max_workers'], + } + + +def load_config(config_data: Dict[str, Any]) -> None: + """Fills in the localhost defaults that the user did not provide""" if 'localhost' not in config_data or not config_data['localhost']: config_data['localhost'] = {} - for key in DEFAULT_CONFIG_KEYS: - if key not in config_data['localhost']: - config_data['localhost'][key] = DEFAULT_CONFIG_KEYS[key] + for key, value in DEFAULT_CONFIG_KEYS.items(): + config_data['localhost'].setdefault(key, value) + # This machine is the only worker, whatever the user configured config_data['localhost']['max_workers'] = 1 - if 'execution_timeout' not in config_data['lithops']: - config_data['lithops']['execution_timeout'] = LOCALHOST_EXECUTION_TIMEOUT - - if 'storage' not in config_data['lithops']: - config_data['lithops']['storage'] = 'localhost' + lithops_cfg = config_data.setdefault('lithops', {}) + lithops_cfg.setdefault('execution_timeout', LOCALHOST_EXECUTION_TIMEOUT) + lithops_cfg.setdefault('storage', 'localhost') diff --git a/lithops/localhost/utils.py b/lithops/localhost/utils.py new file mode 100644 index 000000000..e62f12325 --- /dev/null +++ b/lithops/localhost/utils.py @@ -0,0 +1,188 @@ +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import os +import signal +import shutil +import logging +import tempfile +from typing import Iterable, List, Optional + +try: + import fcntl +except ImportError: # Windows + fcntl = None + +from lithops.constants import LITHOPS_TEMP_DIR + +_COPY_IGNORE = shutil.ignore_patterns('__pycache__', '*.pyc', '*.pyo') + + +def copy_lithops_package( + lithops_location: str, + runner_src: str, + runner_dst: str, + temp_dir: str = LITHOPS_TEMP_DIR, +) -> None: + """ + Copies the Lithops package into the local temp dir and installs the runner. + + Concurrent FunctionExecutor setups share this destination. Copy into a + staging directory first, then replace the destination under a file lock + so one rmtree cannot delete another copy mid-flight. Bytecode caches are + omitted because pytest and other processes rewrite them while we copy. + """ + os.makedirs(temp_dir, exist_ok=True) + dst_path = os.path.join(temp_dir, 'lithops') + lock_path = os.path.join(temp_dir, '.lithops-copy.lock') + staging = tempfile.mkdtemp(prefix='lithops-src-', dir=temp_dir) + try: + shutil.copytree( + lithops_location, + os.path.join(staging, 'lithops'), + ignore=_COPY_IGNORE, + ) + with open(lock_path, 'a') as lock_file: + if fcntl is not None: + fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) + try: + shutil.rmtree(dst_path, ignore_errors=True) + shutil.move(os.path.join(staging, 'lithops'), dst_path) + finally: + if fcntl is not None: + fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) + shutil.copyfile(runner_src, runner_dst) + finally: + shutil.rmtree(staging, ignore_errors=True) + + +def decode_process_output(data: Optional[object]) -> str: + """ + Returns the captured output of a process as text, empty if there is none + """ + if isinstance(data, bytes): + return data.decode('utf-8', errors='replace').strip() + if isinstance(data, str): + return data.strip() + return '' + + +def _read_log_tail(log_file: str, lines: int = 80) -> str: + if not os.path.isfile(log_file): + return '' + try: + with open(log_file, 'r', errors='replace') as fh: + return ''.join(fh.readlines()[-lines:]).strip() + except OSError: + return '' + + +def log_process_failure( + logger: logging.Logger, + message: str, + stdout: Optional[object] = None, + stderr: Optional[object] = None, + log_file: Optional[str] = None, +) -> None: + """ + Logs a localhost worker crash. Reports what the process printed, or the + tail of the runner log when it died without printing anything, as that is + where an import error or a missing dependency shows up + """ + logger.error(message) + detail = decode_process_output(stderr) or decode_process_output(stdout) + if detail: + logger.error(detail) + return + + tail = _read_log_tail(log_file) if log_file else '' + if tail: + logger.error(f'Runner log:\n{tail}') + + +def kill_process(process, is_unix: bool) -> None: + """ + Kills a running subprocess. On Unix the whole process group goes down, so + that the workers the runner forked do not outlive it + """ + if not process or process.poll() is not None: + return + pid = process.pid + if is_unix: + os.killpg(os.getpgid(pid), signal.SIGKILL) + else: + os.kill(pid, signal.SIGTERM) + + +def docker_pull_cmd(docker_path: str, image: str) -> List[str]: + """Builds the command that pulls a runtime image""" + return [docker_path, 'pull', image] + + +def docker_rm_cmd(docker_path: str, name: str) -> List[str]: + """Builds the command that force removes a container""" + return [docker_path, 'rm', '-f', name] + + +def docker_run_cmd( + docker_path: str, + image: str, + *, + name: str, + tmp_path: str, + uid: Optional[int] = None, + gid: Optional[int] = None, + is_podman: bool = False, + use_gpu: bool = False, + extra_run_args: Optional[Iterable[str]] = None, + entrypoint: Optional[str] = 'python3', + container_args: Optional[Iterable[str]] = None, +) -> List[str]: + """ + Builds the command that runs a container with the local temp dir mounted + on /tmp, which is how the runner and the job files reach the container. + + Podman maps the calling user into the container by itself, so --user is + only passed to Docker. + """ + cmd = [docker_path, 'run', '--name', name] + if use_gpu: + cmd.extend(['--gpus', 'all']) + if uid is not None and gid is not None and not is_podman: + cmd.extend(['--user', f'{uid}:{gid}']) + cmd.extend([ + '--env', f'USER={os.getenv("USER", "root")}', + '--rm', '-v', f'{tmp_path}:/tmp', + ]) + if extra_run_args: + cmd.extend(extra_run_args) + if entrypoint is not None: + cmd.extend(['--entrypoint', entrypoint]) + cmd.append(image) + if container_args: + cmd.extend(container_args) + return cmd + + +def docker_exec_python_cmd( + docker_path: str, + container_name: str, + script_path: str, + *script_args: str, +) -> List[str]: + """Builds the command that runs a Python script in a running container""" + inner = ' '.join(['python3', script_path, *script_args]) + return [ + docker_path, 'exec', container_name, '/bin/bash', '-c', inner + ] diff --git a/lithops/localhost/v1/localhost.py b/lithops/localhost/v1/localhost.py index 5d2cd9cfc..ec4a4f116 100644 --- a/lithops/localhost/v1/localhost.py +++ b/lithops/localhost/v1/localhost.py @@ -16,50 +16,63 @@ import os import json -import shlex import queue -import signal import lithops import logging -import shutil import threading import subprocess as sp -from shutil import copyfile +from contextlib import contextmanager from pathlib import Path +from typing import Any, Dict, List -from lithops.version import __version__ from lithops.constants import ( TEMP_DIR, USER_TEMP_DIR, LITHOPS_TEMP_DIR, COMPUTE_CLI_MSG, - JOBS_PREFIX + JOBS_PREFIX, + RN_LOG_FILE, ) from lithops.utils import ( BackendType, get_docker_path, is_lithops_worker, is_podman, - is_unix_system + is_unix_system, + log_prefix, ) from lithops.localhost.config import ( - LocvalhostEnvironment, - get_environment + LocalhostEnvironment, + get_environment, + runtime_info, + runtime_key, +) +from lithops.localhost.utils import ( + copy_lithops_package, + docker_pull_cmd, + docker_rm_cmd, + docker_run_cmd, + kill_process, + log_process_failure, ) logger = logging.getLogger(__name__) RUNNER_FILE = os.path.join(LITHOPS_TEMP_DIR, 'localhost-runner.py') LITHOPS_LOCATION = os.path.dirname(os.path.abspath(lithops.__file__)) +# The local temp dir is mounted on /tmp, so this is where a container sees the +# runner that was copied to RUNNER_FILE +DOCKER_RUNNER_FILE = f'/tmp/{USER_TEMP_DIR}/localhost-runner.py' class LocalhostHandlerV1: """ - A localhostHandler object is used by invokers and other components to access - underlying localhost backend without exposing the implementation details. + A LocalhostHandler object is used by invokers and other components to + access the underlying localhost backend without exposing implementation + details. """ - def __init__(self, config): + def __init__(self, config: Dict[str, Any]): logger.debug('Creating Localhost compute client') self.config = config self.runtime_name = self.config['runtime'] @@ -68,63 +81,96 @@ def __init__(self, config): self.env = None self.job_queue = queue.Queue() self.job_manager = None - self.invocation_in_progress = False + self.invocations_lock = threading.Lock() + self.invocations_in_progress = 0 - msg = COMPUTE_CLI_MSG.format('Localhost compute v1') - logger.info(f"{msg}") + logger.info(COMPUTE_CLI_MSG.format('Localhost compute v1')) def get_backend_type(self): - """ - Wrapper method that returns the type of the backend (Batch or FaaS) - """ + """Returns the backend type, which is invoked with a whole job""" return BackendType.BATCH.value def init(self): - """ - Init tasks for localhost - """ - if self.environment == LocvalhostEnvironment.DEFAULT: + """Creates and sets up the environment where the jobs will run""" + if self.environment == LocalhostEnvironment.DEFAULT: self.env = DefaultEnvironment(self.config) else: self.env = ContainerEnvironment(self.config) - self.env.setup() - def start_manager(self): + @property + def invocation_in_progress(self) -> bool: + """True while at least one invoke() is still queueing its job""" + return self.invocations_in_progress > 0 + + @contextmanager + def _invocation(self): + """ + Marks an invocation as in flight, so that the job manager does not + stop while its job is still being queued + """ + with self.invocations_lock: + self.invocations_in_progress += 1 + try: + yield + finally: + with self.invocations_lock: + self.invocations_in_progress -= 1 + + def _has_pending_work(self) -> bool: + return self.invocation_in_progress or not self.job_queue.empty() + + def _run_queued_job( + self, job_payload: Dict[str, Any], job_filename: str + ) -> None: """ - Starts manager thread to keep order in tasks + Runs one job in the localhost worker and waits until it finishes, so + that queued jobs run one after another """ + executor_id = job_payload['executor_id'] + job_id = job_payload['job_id'] + logger.debug( + f'{log_prefix(executor_id, job_id)} - Running ' + f'{len(job_payload["call_ids"])} activations in the localhost worker' + ) + process = self.env.run_job(job_payload['job_key'], job_filename) + if process is None: + logger.debug( + f'{log_prefix(executor_id, job_id)} - Job was stopped ' + 'before starting' + ) + return + stdout, stderr = process.communicate() + + if process.returncode != 0: + log_process_failure( + logger, + f'{log_prefix(executor_id, job_id)} - Job process failed ' + f'with return code {process.returncode}', + stdout=stdout, + stderr=stderr, + log_file=RN_LOG_FILE, + ) + logger.debug(f'{log_prefix(executor_id, job_id)} - Execution finished') + + def start_manager(self): + """ + Starts the thread that drains the job queue, unless it already runs + """ def job_manager(): - logger.debug('Staring localhost job manager') + logger.debug('Starting localhost job manager') while True: job_payload, job_filename = self.job_queue.get() + is_sentinel = job_payload is None and job_filename is None + if not is_sentinel: + self._run_queued_job(job_payload, job_filename) - if job_payload is None and job_filename is None: - if self.invocation_in_progress or not self.job_queue.empty(): - continue - else: - break - - executor_id = job_payload['executor_id'] - job_id = job_payload['job_id'] - total_calls = len(job_payload['call_ids']) - job_key = job_payload['job_key'] - logger.debug(f'ExecutorID {executor_id} | JobID {job_id} - Running ' - f'{total_calls} activations in the localhost worker') - process = self.env.run_job(job_key, job_filename) - process.communicate() # blocks until the process finishes - if process.returncode != 0: - logger.error(f"ExecutorID {executor_id} | JobID {job_id} - Job " - f"process failed with return code {process.returncode}") - logger.debug(f'ExecutorID {executor_id} | JobID {job_id} - Execution finished') - - if self.job_queue.empty(): - if self.invocation_in_progress: - continue - else: - break + # An invocation in flight is about to queue its job, so the + # manager only stops once there is nothing left to run + if not self._has_pending_work(): + break self.job_manager = None logger.debug("Localhost job manager finished") @@ -134,189 +180,214 @@ def job_manager(): self.job_manager.start() def deploy_runtime(self, runtime_name, *args): - """ - Extract the runtime metadata and preinstalled modules - """ + """Returns the metadata of the runtime, which needs no deployment""" logger.info(f"Deploying runtime: {runtime_name}") return self.env.get_metadata() - def invoke(self, job_payload): - """ - Run the job description against the selected environment - """ - self.invocation_in_progress = True - executor_id = job_payload['executor_id'] - job_id = job_payload['job_id'] - - logger.debug(f'ExecutorID {executor_id} | JobID {job_id} - Putting job into localhost queue') - - job_filename = self.env.prepare_job_file(job_payload) - self.job_queue.put((job_payload, job_filename)) - - self.start_manager() - self.invocation_in_progress = False + def invoke(self, job_payload: Dict[str, Any]) -> None: + """Queues a job and makes sure that the job manager is running""" + with self._invocation(): + executor_id = job_payload['executor_id'] + job_id = job_payload['job_id'] + logger.debug( + f'{log_prefix(executor_id, job_id)} - ' + 'Putting job into localhost queue' + ) + job_filename = self.env.prepare_job_file(job_payload) + self.job_queue.put((job_payload, job_filename)) + self.start_manager() def get_runtime_key(self, runtime_name, *args): - """ - Generate the runtime key that identifies the runtime - """ - runtime_key = os.path.join('localhost', __version__, runtime_name.strip("/")) - - return runtime_key + """Returns the key the runtime metadata is cached under""" + return runtime_key(runtime_name) def get_runtime_info(self): - """ - Method that returns a dictionary with all the relevant runtime information - set in config - """ - return { - 'runtime_name': self.config['runtime'], - 'runtime_memory': self.config.get('runtime_memory'), - 'runtime_timeout': self.config.get('runtime_timeout'), - 'max_workers': self.config['max_workers'], - } + """Returns the runtime limits the executor reports to the user""" + return runtime_info(self.config) def clean(self, **kwargs): - """ - Deletes all local runtimes - """ + """Nothing to clean up: the localhost backend deploys nothing""" pass def clear(self, job_keys=None, exception=None): """ - Kills all running jobs processes + Drops the given jobs if they have not started yet and kills them if + they have. Jobs that were not named stay queued """ - while not self.job_queue.empty(): - try: - self.job_queue.get(False) - except Exception: - pass - + self._drop_queued_jobs(job_keys) self.env.stop(job_keys) if self.job_manager: self.job_queue.put((None, None)) + def _drop_queued_jobs(self, job_keys=None) -> None: + """ + Takes the given jobs out of the queue, putting back the ones that + belong to jobs nobody asked to stop + """ + kept = [] + while True: + try: + queued_job = self.job_queue.get(block=False) + except queue.Empty: + break + job_payload, _ = queued_job + # A sentinel only tells the manager to look at its work again, + # and clear() queues a fresh one right after this + if job_payload is None or job_keys is None: + continue + if job_payload['job_key'] not in job_keys: + kept.append(queued_job) + + for queued_job in kept: + self.job_queue.put(queued_job) + class ExecutionEnvironment: - """ - Base environment class for shared methods - """ + """Base environment class for shared methods.""" - def __init__(self, config): + def __init__(self, config: Dict[str, Any]): self.config = config self.runtime_name = self.config['runtime'] self.is_unix_system = is_unix_system() - self.jobs = {} # dict to store executed jobs (job_keys) and PIDs + self.jobs = {} + self.jobs_lock = threading.Lock() + self.stopped_jobs = set() def _copy_lithops_to_tmp(self): + # A job invoked from inside a worker reuses the package that the + # parent already copied, otherwise it would overwrite it while running if is_lithops_worker() and os.path.isfile(RUNNER_FILE): return - os.makedirs(LITHOPS_TEMP_DIR, exist_ok=True) - dst_path = os.path.join(LITHOPS_TEMP_DIR, 'lithops') - shutil.rmtree(dst_path, ignore_errors=True) - shutil.copytree(LITHOPS_LOCATION, dst_path, dirs_exist_ok=True) - src_handler = os.path.join(LITHOPS_LOCATION, 'localhost', 'v1', 'runner.py') - copyfile(src_handler, RUNNER_FILE) - - def prepare_job_file(self, job_payload): + copy_lithops_package( + LITHOPS_LOCATION, + os.path.join(LITHOPS_LOCATION, 'localhost', 'v1', 'runner.py'), + RUNNER_FILE, + LITHOPS_TEMP_DIR, + ) + + def _ensure_runner(self): + if not os.path.isfile(RUNNER_FILE): + self.setup() + + def prepare_job_file(self, job_payload: Dict[str, Any]) -> str: """ - Creates the job file that contains the job payload to be executed + Dumps the job payload where the runner will read it, and returns the + path as the runner sees it """ job_key = job_payload['job_key'] + with self.jobs_lock: + self.stopped_jobs.discard(job_key) storage_backend = job_payload['config']['lithops']['storage'] storage_bucket = job_payload['config'][storage_backend]['storage_bucket'] - local_job_dir = os.path.join(LITHOPS_TEMP_DIR, storage_bucket, JOBS_PREFIX) + local_job_dir = os.path.join( + LITHOPS_TEMP_DIR, storage_bucket, JOBS_PREFIX + ) docker_job_dir = f'/tmp/{USER_TEMP_DIR}/{storage_bucket}/{JOBS_PREFIX}' job_file = f'{job_key}-job.json' os.makedirs(local_job_dir, exist_ok=True) local_job_filename = os.path.join(local_job_dir, job_file) - with open(local_job_filename, 'w') as jl: - json.dump(job_payload, jl, default=str) - - if isinstance(self, ContainerEnvironment): - job_filename = f'{docker_job_dir}/{job_file}' - else: - job_filename = local_job_filename + with open(local_job_filename, 'w') as job_file_handle: + json.dump(job_payload, job_file_handle, default=str) - return job_filename + return self._job_file_for_runner( + local_job_filename, f'{docker_job_dir}/{job_file}' + ) - def stop(self, job_keys=None): + def _job_file_for_runner(self, local_path: str, container_path: str) -> str: """ - Stops running processes + Returns the path the runner reads the job file from, which by default + is where it was written """ + return local_path - def kill_job(job_key): - if self.jobs[job_key].poll() is None: - logger.debug(f'Killing job {job_key} with PID {self.jobs[job_key].pid}') - PID = self.jobs[job_key].pid - if self.is_unix_system: - PGID = os.getpgid(PID) - os.killpg(PGID, signal.SIGKILL) - else: - os.kill(PID, signal.SIGTERM) - del self.jobs[job_key] + def _start_job_process(self, job_key: str, cmd: List[str]): + """ + Starts the process that runs a whole job and registers it as one step, + so that a stop() running right now either kills this process or keeps + it from being started at all. Returns None in the latter case + """ + with self.jobs_lock: + if job_key in self.stopped_jobs: + logger.debug(f'Job {job_key} not started, it was stopped') + return None + process = sp.Popen( + cmd, + stdout=sp.PIPE, + stderr=sp.PIPE, + start_new_session=True, + ) + self.jobs[job_key] = process + return process + def stop(self, job_keys=None): + """Kills the job processes that are still running""" to_delete = job_keys or list(self.jobs.keys()) - for job_key in to_delete: - try: - if job_key in self.jobs: - kill_job(job_key) - except Exception: - pass + with self.jobs_lock: + # Marked before the sweep, so that a job about to start is not + # left running behind it + self.stopped_jobs.update(to_delete) + for job_key in to_delete: + try: + if job_key not in self.jobs: + continue + process = self.jobs[job_key] + logger.debug( + f'Killing job {job_key} with PID {process.pid}' + ) + kill_process(process, self.is_unix_system) + del self.jobs[job_key] + except Exception: + pass class DefaultEnvironment(ExecutionEnvironment): - """ - Default environment uses current python3 installation - """ + """Default environment uses the current Python installation.""" - def __init__(self, config): + def __init__(self, config: Dict[str, Any]): super().__init__(config) logger.debug(f'Starting default environment for {self.runtime_name}') def setup(self): + """Installs the Lithops package and the runner in the temp dir""" logger.debug('Setting up default environment') self._copy_lithops_to_tmp() def get_metadata(self): - if not os.path.isfile(RUNNER_FILE): - self.setup() + """Asks the local interpreter for the packages it has installed""" + self._ensure_runner() logger.debug(f"Extracting metadata from: {self.runtime_name}") - cmd = [self.runtime_name, RUNNER_FILE, 'get_metadata'] process = sp.run( - cmd, check=True, + [self.runtime_name, RUNNER_FILE, 'get_metadata'], + check=True, stdout=sp.PIPE, universal_newlines=True, - start_new_session=True + start_new_session=True, ) - runtime_meta = json.loads(process.stdout.strip()) - return runtime_meta + return json.loads(process.stdout.strip()) - def run_job(self, job_key, job_filename): - """ - Runs a job - """ - if not os.path.isfile(RUNNER_FILE): - self.setup() - - cmd = [self.runtime_name, RUNNER_FILE, 'run_job', job_filename] - process = sp.Popen(cmd, stdout=sp.PIPE, stderr=sp.PIPE, start_new_session=True) - self.jobs[job_key] = process + def run_job(self, job_key: str, job_filename: str): + """Starts the runner that executes the whole job, and returns it""" + self._ensure_runner() - return process + return self._start_job_process( + job_key, + [self.runtime_name, RUNNER_FILE, 'run_job', job_filename], + ) class ContainerEnvironment(ExecutionEnvironment): - """ - Docker environment uses a docker runtime image - """ + """Container environment uses a container runtime image.""" + + def _job_file_for_runner(self, local_path: str, container_path: str) -> str: + """The container sees the local temp dir mounted on /tmp""" + return container_path - def __init__(self, config): + def __init__(self, config: Dict[str, Any]): super().__init__(config) logger.debug(f'Starting container environment for {self.runtime_name}') self.use_gpu = self.config.get('use_gpu', False) @@ -325,69 +396,70 @@ def __init__(self, config): self.uid = os.getuid() if self.is_unix_system else None self.gid = os.getgid() if self.is_unix_system else None + def _container_cmd(self, name, container_args, use_gpu=False) -> List[str]: + return docker_run_cmd( + self.docker_path, + self.runtime_name, + name=name, + tmp_path=Path(TEMP_DIR).as_posix(), + uid=self.uid, + gid=self.gid, + is_podman=self.is_podman, + use_gpu=use_gpu, + container_args=container_args, + ) + def setup(self): + """Installs the runner in the temp dir and pulls the image if asked""" logger.debug('Setting up container environment') self._copy_lithops_to_tmp() if self.config.get('pull_runtime', False): logger.debug(f'Pulling runtime {self.runtime_name}') sp.run( - shlex.split(f'docker pull {self.runtime_name}'), check=True, - stdout=sp.PIPE, universal_newlines=True + docker_pull_cmd(self.docker_path, self.runtime_name), + check=True, + stdout=sp.PIPE, + universal_newlines=True, ) def get_metadata(self): - if not os.path.isfile(RUNNER_FILE): - self.setup() + """Asks the runtime image for the packages it has installed""" + self._ensure_runner() logger.debug(f"Extracting metadata from: {self.runtime_name}") - - tmp_path = Path(TEMP_DIR).as_posix() - - cmd = f'{self.docker_path} run --name lithops_metadata ' - cmd += f'--user {self.uid}:{self.gid} ' if self.is_unix_system and not self.is_podman else '' - cmd += f'--env USER={os.getenv("USER", "root")} ' - cmd += f'--rm -v {tmp_path}:/tmp --entrypoint "python3" ' - cmd += f'{self.runtime_name} /tmp/{USER_TEMP_DIR}/localhost-runner.py get_metadata' - process = sp.run( - shlex.split(cmd), check=True, stdout=sp.PIPE, - universal_newlines=True, start_new_session=True + self._container_cmd( + 'lithops_metadata', [DOCKER_RUNNER_FILE, 'get_metadata'] + ), + check=True, + stdout=sp.PIPE, + universal_newlines=True, + start_new_session=True, ) - runtime_meta = json.loads(process.stdout.strip()) + return json.loads(process.stdout.strip()) - return runtime_meta - - def run_job(self, job_key, job_filename): + def run_job(self, job_key: str, job_filename: str): """ - Runs a job + Starts a container that executes the whole job, and returns the + process that runs it """ - if not os.path.isfile(RUNNER_FILE): - self.setup() - - tmp_path = Path(TEMP_DIR).as_posix() - - cmd = f'{self.docker_path} run --name lithops_{job_key} ' - cmd += '--gpus all ' if self.use_gpu else '' - cmd += f'--user {self.uid}:{self.gid} ' if self.is_unix_system and not self.is_podman else '' - cmd += f'--env USER={os.getenv("USER", "root")} ' - cmd += f'--rm -v {tmp_path}:/tmp --entrypoint "python3" ' - cmd += f'{self.runtime_name} /tmp/{USER_TEMP_DIR}/localhost-runner.py run_job {job_filename}' - - process = sp.Popen(shlex.split(cmd), stdout=sp.PIPE, stderr=sp.PIPE, start_new_session=True) - self.jobs[job_key] = process - - return process + self._ensure_runner() + + return self._start_job_process( + job_key, + self._container_cmd( + f'lithops_{job_key}', + [DOCKER_RUNNER_FILE, 'run_job', job_filename], + use_gpu=self.use_gpu, + ), + ) def stop(self, job_keys=None): - """ - Stops running containers - """ - jk_to_delete = job_keys or list(self.jobs.keys()) - - for job_key in jk_to_delete: + """Removes the job containers and kills the processes behind them""" + for job_key in job_keys or list(self.jobs.keys()): sp.Popen( - shlex.split(f'{self.docker_path} rm -f lithops_{job_key}'), - stdout=sp.DEVNULL, stderr=sp.DEVNULL + docker_rm_cmd(self.docker_path, f'lithops_{job_key}'), + stdout=sp.DEVNULL, + stderr=sp.DEVNULL, ) - super().stop(job_keys) diff --git a/lithops/localhost/v1/runner.py b/lithops/localhost/v1/runner.py index 7ac39b774..1126f1cb8 100644 --- a/lithops/localhost/v1/runner.py +++ b/lithops/localhost/v1/runner.py @@ -20,49 +20,74 @@ import platform import logging import uuid +import traceback import multiprocessing as mp from pathlib import Path from lithops.worker import function_handler from lithops.worker.utils import get_runtime_metadata -from lithops.constants import LITHOPS_TEMP_DIR, JOBS_DIR, LOGS_DIR, \ - RN_LOG_FILE, LOGGER_FORMAT +from lithops.utils import log_prefix +from lithops.constants import ( + LITHOPS_TEMP_DIR, + JOBS_DIR, + LOGS_DIR, + RN_LOG_FILE, + LOGGER_FORMAT, +) -log_file_stream = open(RN_LOG_FILE, 'a') - -os.makedirs(LITHOPS_TEMP_DIR, exist_ok=True) -os.makedirs(JOBS_DIR, exist_ok=True) -os.makedirs(LOGS_DIR, exist_ok=True) - -logging.basicConfig(stream=log_file_stream, - level=logging.INFO, - format=LOGGER_FORMAT) logger = logging.getLogger('lithops.localhost.runner') -# Python 3.14 defaults to forkserver on Linux; Lithops requires fork. -if platform.system() != 'Windows': +def _configure_runner_logging(): + """ + Sends the runner logs to the runner log file, and returns the stream so + that the caller can also redirect the job output to it + """ + os.makedirs(LITHOPS_TEMP_DIR, exist_ok=True) + os.makedirs(JOBS_DIR, exist_ok=True) + os.makedirs(LOGS_DIR, exist_ok=True) + log_file_stream = open(RN_LOG_FILE, 'a') + logging.basicConfig( + stream=log_file_stream, + level=logging.INFO, + format=LOGGER_FORMAT, + ) + return log_file_stream + + +def _set_fork_start_method(): + """ + Forces fork, which Lithops relies on for its workers to inherit the job. + Python 3.14 defaults to forkserver on Linux + """ + if platform.system() == 'Windows': + return try: mp.set_start_method('fork') except RuntimeError: + # Already set by an earlier call in this interpreter pass -def run_job(): +def run_job(log_file_stream): + """ + Runs the whole job described by the job file given as the second argument + """ + # This process has no console: anything printed goes to the runner log sys.stdout = log_file_stream sys.stderr = log_file_stream job_filename = sys.argv[2] logger.info(f'Got {job_filename} job file') - with open(job_filename, 'rb') as jf: - job_payload = json.load(jf) + with open(job_filename, 'r') as job_file: + job_payload = json.load(job_file) executor_id = job_payload['executor_id'] job_id = job_payload['job_id'] job_key = job_payload['job_key'] - logger.info(f'ExecutorID {executor_id} | JobID {job_id} - Starting execution') + logger.info(f'{log_prefix(executor_id, job_id)} - Starting execution') act_id = str(uuid.uuid4()).replace('-', '')[:12] os.environ['__LITHOPS_ACTIVATION_ID'] = act_id @@ -73,29 +98,44 @@ def run_job(): except KeyboardInterrupt: pass - done = os.path.join(JOBS_DIR, job_key + '.done') - Path(done).touch() + Path(os.path.join(JOBS_DIR, job_key + '.done')).touch() if os.path.exists(job_filename): os.remove(job_filename) - logger.info(f'ExecutorID {executor_id} | JobID {job_id} - Execution Finished') + logger.info(f'{log_prefix(executor_id, job_id)} - Execution Finished') def extract_runtime_meta(): - runtime_meta = get_runtime_metadata() - print(json.dumps(runtime_meta)) + """Prints the metadata of this runtime, which the client reads back""" + print(json.dumps(get_runtime_metadata())) -if __name__ == "__main__": - logger.info('Starting Localhost job runner') - command = sys.argv[1] - logger.info(f'Received command: {command}') +def main(): + """Entry point of the runner subprocess, dispatching the argv command""" + _set_fork_start_method() + log_file_stream = _configure_runner_logging() + try: + logger.info('Starting Localhost job runner') + command = sys.argv[1] + logger.info(f'Received command: {command}') + + handlers = { + 'get_metadata': extract_runtime_meta, + 'run_job': lambda: run_job(log_file_stream), + } + handler = handlers.get(command) + if handler is None: + logger.error(f'Invalid command: {command}') + sys.exit(1) + handler() + except Exception: + logger.exception('Localhost job runner failed') + traceback.print_exc(file=sys.__stderr__) + sys.exit(1) + finally: + log_file_stream.close() - switcher = { - 'get_metadata': extract_runtime_meta, - 'run_job': run_job - } - switcher.get(command, lambda: "Invalid command")() - log_file_stream.close() +if __name__ == "__main__": + main() diff --git a/lithops/localhost/v2/localhost.py b/lithops/localhost/v2/localhost.py index a873104af..87612b1ca 100644 --- a/lithops/localhost/v2/localhost.py +++ b/lithops/localhost/v2/localhost.py @@ -18,18 +18,16 @@ import os import json import threading +import time import uuid -import shlex -import signal import lithops import logging -import shutil import queue import subprocess as sp -from shutil import copyfile +from contextlib import contextmanager from pathlib import Path +from typing import Any, Dict, List -from lithops.version import __version__ from lithops.constants import ( JOBS_DIR, TEMP_DIR, @@ -37,6 +35,7 @@ COMPUTE_CLI_MSG, CPU_COUNT, USER_TEMP_DIR, + RN_LOG_FILE, ) from lithops.utils import ( BackendType, @@ -44,71 +43,108 @@ get_docker_path, is_lithops_worker, is_podman, - is_unix_system + is_unix_system, + log_prefix, ) from lithops.localhost.config import ( - LocvalhostEnvironment, - get_environment + LocalhostEnvironment, + get_environment, + runtime_info, + runtime_key, +) +from lithops.localhost.utils import ( + copy_lithops_package, + docker_exec_python_cmd, + docker_pull_cmd, + docker_rm_cmd, + docker_run_cmd, + kill_process, + log_process_failure, ) logger = logging.getLogger(__name__) RUNNER_FILE = os.path.join(LITHOPS_TEMP_DIR, 'localhost-runner.py') LITHOPS_LOCATION = os.path.dirname(os.path.abspath(lithops.__file__)) +# The local temp dir is mounted on /tmp, so this is where a container sees the +# runner that was copied to RUNNER_FILE +DOCKER_RUNNER_FILE = f'/tmp/{USER_TEMP_DIR}/localhost-runner.py' +# How long the job manager waits before looking again for the latch of +# a job that is being invoked right now +MANAGER_IDLE_WAIT = 0.1 class LocalhostHandlerV2: """ - A localhostHandler object is used by invokers and other components to - access underlying localhost backend without exposing the implementation + A LocalhostHandler object is used by invokers and other components to + access the underlying localhost backend without exposing implementation details. """ - def __init__(self, localhost_config): + def __init__(self, config: Dict[str, Any]): logger.debug('Creating Localhost compute client') - self.config = localhost_config + self.config = config self.runtime_name = self.config['runtime'] self.environment = get_environment(self.runtime_name) self.env = None self.job_manager = None - self.invocation_in_progress = False + self.invocations_lock = threading.Lock() + self.invocations_in_progress = 0 - msg = COMPUTE_CLI_MSG.format('Localhost compute v2') - logger.info(f"{msg}") + logger.info(COMPUTE_CLI_MSG.format('Localhost compute v2')) - def get_backend_type(self): + @property + def invocation_in_progress(self) -> bool: + """True while at least one invoke() is still queueing its tasks""" + return self.invocations_in_progress > 0 + + @contextmanager + def _invocation(self): """ - Wrapper method that returns the type of the backend (Batch or FaaS) + Marks an invocation as in flight, so that the job manager does not + stop while its tasks are still being queued """ + with self.invocations_lock: + self.invocations_in_progress += 1 + try: + yield + finally: + with self.invocations_lock: + self.invocations_in_progress -= 1 + + def get_backend_type(self): + """Returns the backend type, which is invoked with a whole job""" return BackendType.BATCH.value def init(self): - """ - Init tasks for localhost - """ - if self.environment == LocvalhostEnvironment.DEFAULT: + """Creates and sets up the environment where the tasks will run""" + if self.environment == LocalhostEnvironment.DEFAULT: self.env = DefaultEnvironment(self.config) else: self.env = ContainerEnvironment(self.config) - self.env.setup() def start_manager(self): """ - Starts manager thread to keep order in tasks + Starts the thread that waits for the running jobs, and the consumers + that execute their tasks, unless they are already running """ def job_manager(): - logger.debug('Staring localhost job manager') + logger.debug('Starting localhost job manager') while True: for job_key in list(self.env.jobs.keys()): self.env.jobs[job_key].wait() - if all(job.done for job in self.env.jobs.values()): + # A new job may have been invoked while waiting for the + # previous ones, so only stop once every latch is down + if all(job.done for job in list(self.env.jobs.values())): if self.invocation_in_progress: + # An invoke() is queueing its tasks right now, so + # wait for its latch to show up instead of spinning + time.sleep(MANAGER_IDLE_WAIT) continue - else: - break + break self.job_manager = None logger.debug("Localhost job manager finished") @@ -119,250 +155,298 @@ def job_manager(): self.env.start() def deploy_runtime(self, runtime_name, *args): - """ - Extract the runtime metadata and preinstalled modules - """ + """Returns the metadata of the runtime, which needs no deployment""" logger.info(f"Deploying runtime: {runtime_name}") return self.env.get_metadata() - def invoke(self, job_payload): - """ - Run the job description against the selected environment - """ - self.invocation_in_progress = True - executor_id = job_payload['executor_id'] - job_id = job_payload['job_id'] - total_calls = len(job_payload['call_ids']) - - logger.debug(f'ExecutorID {executor_id} | JobID {job_id} - Running ' - f'{total_calls} activations in the localhost worker') - - self.env.run_job(job_payload) - - self.start_manager() - self.invocation_in_progress = False + def invoke(self, job_payload: Dict[str, Any]) -> None: + """Queues the tasks of a job and makes sure the consumers are up""" + with self._invocation(): + executor_id = job_payload['executor_id'] + job_id = job_payload['job_id'] + logger.debug( + f'{log_prefix(executor_id, job_id)} - Running ' + f'{len(job_payload["call_ids"])} activations in the localhost ' + f'worker' + ) + self.env.run_job(job_payload) + self.start_manager() def get_runtime_key(self, runtime_name, *args): - """ - Generate the runtime key that identifies the runtime - """ - runtime_key = os.path.join('localhost', __version__, runtime_name.strip("/")) - - return runtime_key + """Returns the key the runtime metadata is cached under""" + return runtime_key(runtime_name) def get_runtime_info(self): - """ - Method that returns a dictionary with all the relevant runtime - information set in config - """ - return { - 'runtime_name': self.config['runtime'], - 'runtime_memory': self.config.get('runtime_memory'), - 'runtime_timeout': self.config.get('runtime_timeout'), - 'max_workers': self.config['max_workers'], - } + """Returns the runtime limits the executor reports to the user""" + return runtime_info(self.config) def clean(self, **kwargs): - """ - Deletes all local runtimes - """ + """Nothing to clean up: the localhost backend deploys nothing""" pass def clear(self, job_keys=None, exception=None): """ - Kills the running service in case of exception + Drops the tasks of the given jobs that have not started yet, kills the + running ones and releases their latches so that the job manager can + finish. Jobs that were not named are left running """ - while not self.env.work_queue.empty(): - try: - self.env.work_queue.get(False) - except Exception: - pass - + self.env.drop_pending_tasks(job_keys) self.env.stop(job_keys) for job_key in list(self.env.jobs.keys()): + if job_keys is not None and job_key not in job_keys: + continue while not self.env.jobs[job_key].done: self.env.jobs[job_key].unlock() class ExecutionEnvironment: - """ - Base environment class for shared methods - """ + """Base environment class for shared methods.""" - def __init__(self, config): + def __init__(self, config: Dict[str, Any]): self.config = config self.runtime_name = self.config['runtime'] self.worker_processes = self.config.get('worker_processes', CPU_COUNT) self.work_queue = queue.Queue() self.is_unix_system = is_unix_system() self.task_processes = {} + self.task_processes_lock = threading.Lock() + self.stopped_jobs = set() self.consumer_threads = [] self.jobs = {} def _copy_lithops_to_tmp(self): + # A task invoked from inside a worker reuses the package that the + # parent already copied, otherwise it would overwrite it while running if is_lithops_worker() and os.path.isfile(RUNNER_FILE): return - os.makedirs(LITHOPS_TEMP_DIR, exist_ok=True) - dst_path = os.path.join(LITHOPS_TEMP_DIR, 'lithops') - shutil.rmtree(dst_path, ignore_errors=True) - shutil.copytree(LITHOPS_LOCATION, dst_path, dirs_exist_ok=True) - src_handler = os.path.join(LITHOPS_LOCATION, 'localhost', 'v2', 'runner.py') - copyfile(src_handler, RUNNER_FILE) - - def run_job(self, job_payload): + copy_lithops_package( + LITHOPS_LOCATION, + os.path.join(LITHOPS_LOCATION, 'localhost', 'v2', 'runner.py'), + RUNNER_FILE, + LITHOPS_TEMP_DIR, + ) + + def _ensure_runner(self): + if not os.path.isfile(RUNNER_FILE): + self.setup() + + def _run_task_process(self, job_key_call_id: str, cmd: List[str]) -> None: + """ + Runs one task in a subprocess and waits for it, reporting whatever it + printed if it failed + """ + logger.debug(f"Going to execute task process {job_key_call_id}") + # Started and registered as one step, so that a stop() running right + # now either kills this process or stops it from being started at all + with self.task_processes_lock: + if job_key_call_id.rsplit('-', 1)[0] in self.stopped_jobs: + logger.debug( + f"Task process {job_key_call_id} not started, its job " + f"was stopped" + ) + return + process = sp.Popen( + cmd, + stdout=sp.PIPE, + stderr=sp.PIPE, + start_new_session=True, + ) + self.task_processes[job_key_call_id] = process + + stdout, stderr = process.communicate() + + if process.returncode != 0: + log_process_failure( + logger, + f"Task process {job_key_call_id} failed with return " + f"code {process.returncode}", + stdout=stdout, + stderr=stderr, + log_file=RN_LOG_FILE, + ) + self.task_processes.pop(job_key_call_id, None) + logger.debug(f"Task process {job_key_call_id} finished") + + def run_job(self, job_payload: Dict[str, Any]) -> None: """ - Adds a job to the localhost work queue + Splits a job into one queued task per call, each one carrying only the + data range of its own call """ job_key = job_payload['job_key'] self.jobs[job_key] = CountDownLatch(len(job_payload['call_ids'])) os.makedirs(os.path.join(JOBS_DIR, job_key), exist_ok=True) dbr = job_payload['data_byte_ranges'] + with self.task_processes_lock: + self.stopped_jobs.discard(job_payload['job_key']) for call_id in job_payload['call_ids']: task_payload = copy.deepcopy(job_payload) task_payload['call_ids'] = [call_id] task_payload['data_byte_ranges'] = [dbr[int(call_id)]] self.work_queue.put(json.dumps(task_payload)) - def start(self): + def _process_task(self, task_payload_str: str) -> None: """ - Starts the threads responsible to consume individual tasks from the queue - and execute them in the appropiate environment + Dumps a queued task where the runner will read it, runs it and counts + it down on its job latch """ - if self.consumer_threads: - return + task_payload = json.loads(task_payload_str) + job_key = task_payload['job_key'] + call_id = task_payload['call_ids'][0] - def process_task(task_payload_str): - task_payload = json.loads(task_payload_str) - job_key = task_payload['job_key'] - call_id = task_payload['call_ids'][0] + task_filename = os.path.join(JOBS_DIR, job_key, call_id + '.task') + with open(task_filename, 'w') as task_file: + json.dump(task_payload, task_file, default=str) - task_filename = os.path.join(JOBS_DIR, job_key, call_id + '.task') - with open(task_filename, 'w') as jl: - json.dump(task_payload, jl, default=str) + self.run_task(job_key, call_id) - self.run_task(job_key, call_id) + if os.path.exists(task_filename): + os.remove(task_filename) - if os.path.exists(task_filename): - os.remove(task_filename) + self.jobs[job_key].unlock() - self.jobs[job_key].unlock() + def _queue_consumer(self) -> None: + while True: + task_payload_str = self.work_queue.get() + if task_payload_str is None: + break + self._process_task(task_payload_str) - def queue_consumer(work_queue): - while True: - task_payload_str = work_queue.get() - if task_payload_str is None: - break - process_task(task_payload_str) + def start(self): + """Starts the consumer threads that run the queued tasks""" + if self.consumer_threads: + return logger.debug("Starting Localhost work queue consumer threads") for _ in range(self.worker_processes): - t = threading.Thread( - target=queue_consumer, - args=(self.work_queue,), - daemon=True) - t.start() - self.consumer_threads.append(t) + thread = threading.Thread(target=self._queue_consumer, daemon=True) + thread.start() + self.consumer_threads.append(thread) + + def drop_pending_tasks(self, job_keys=None) -> None: + """ + Takes the queued tasks of the given jobs out of the work queue, and + puts back the ones belonging to jobs nobody asked to stop + """ + kept = [] + while True: + try: + task_payload_str = self.work_queue.get(block=False) + except queue.Empty: + break + # A sentinel left behind by an earlier stop() would kill the next + # consumer that starts, so it never goes back into the queue + if task_payload_str is None or job_keys is None: + continue + if json.loads(task_payload_str)['job_key'] not in job_keys: + kept.append(task_payload_str) + + for task_payload_str in kept: + self.work_queue.put(task_payload_str) def stop(self, job_keys=None): """ - Stops running consumer threads + Kills the task processes of the given jobs, and stops the environment + unless jobs other than those are still to run """ + self._kill_task_processes(job_keys or list(self.jobs.keys())) + + if job_keys is not None and self._has_jobs_left(job_keys): + logger.debug( + "Localhost environment left running, it still has jobs to run" + ) + return + + self._teardown() + + def _has_jobs_left(self, stopped_job_keys) -> bool: + """Tells whether a job other than the stopped ones is still running""" + return any( + not latch.done + for job_key, latch in list(self.jobs.items()) + if job_key not in stopped_job_keys + ) + + def _kill_task_processes(self, job_keys) -> None: + """Kills the task processes of the given jobs""" + with self.task_processes_lock: + # Marked before the sweep, so that a task about to start is not + # left running behind it + self.stopped_jobs.update(job_keys) + for job_key in job_keys: + for job_key_call_id in list(self.task_processes.keys()): + if job_key_call_id.rsplit('-', 1)[0] != job_key: + continue + process = self.task_processes.pop(job_key_call_id, None) + if process is None: + continue + try: + kill_process(process, self.is_unix_system) + except Exception: + pass + + def _teardown(self) -> None: + """Stops the consumer threads, leaving the environment idle""" + if not self.consumer_threads: + return + logger.debug("Stopping Localhost work queue consumer threads") - for _ in range(self.worker_processes): + # One sentinel per running consumer, no more: one that nobody takes + # stays in the queue and kills the next consumer that starts + for _ in self.consumer_threads: self.work_queue.put(None) - for t in self.consumer_threads: - t.join() + for thread in self.consumer_threads: + thread.join() self.consumer_threads = [] class DefaultEnvironment(ExecutionEnvironment): - """ - Default environment uses current python3 installation - """ + """Default environment uses the current Python installation.""" - def __init__(self, config): + def __init__(self, config: Dict[str, Any]): super().__init__(config) logger.debug(f'Starting default environment for {self.runtime_name}') def setup(self): + """Installs the Lithops package and the runner in the temp dir""" logger.debug('Setting up default environment') self._copy_lithops_to_tmp() def get_metadata(self): - if not os.path.isfile(RUNNER_FILE): - self.setup() + """Asks the local interpreter for the packages it has installed""" + self._ensure_runner() logger.debug(f"Extracting metadata from: {self.runtime_name}") - cmd = [self.runtime_name, RUNNER_FILE, 'get_metadata'] process = sp.run( - cmd, check=True, + [self.runtime_name, RUNNER_FILE, 'get_metadata'], + check=True, stdout=sp.PIPE, universal_newlines=True, - start_new_session=True + start_new_session=True, ) - runtime_meta = json.loads(process.stdout.strip()) - return runtime_meta + return json.loads(process.stdout.strip()) def start(self): - if not os.path.isfile(RUNNER_FILE): - self.setup() - + """Starts the consumer threads, with the runner in place""" + self._ensure_runner() super().start() - def run_task(self, job_key, call_id): - """ - Runs a task - """ - job_key_call_id = f'{job_key}-{call_id}' + def run_task(self, job_key: str, call_id: str) -> None: + """Runs one task in a subprocess of the local interpreter""" task_filename = os.path.join(JOBS_DIR, job_key, call_id + '.task') - - logger.debug(f"Going to execute task process {job_key_call_id}") - cmd = [self.runtime_name, RUNNER_FILE, 'run_job', task_filename] - process = sp.Popen(cmd, stdout=sp.PIPE, stderr=sp.PIPE, start_new_session=True) - self.task_processes[job_key_call_id] = process - process.communicate() # blocks until the process finishes - if process.returncode != 0: - logger.error(f"Task process {job_key_call_id} failed with return code {process.returncode}") - del self.task_processes[job_key_call_id] - logger.debug(f"Task process {job_key_call_id} finished") - - def stop(self, job_keys=None): - """ - Stops running processes - """ - def kill_process(process): - if process and process.poll() is None: - PID = process.pid - if self.is_unix_system: - PGID = os.getpgid(PID) - os.killpg(PGID, signal.SIGKILL) - else: - os.kill(PID, signal.SIGTERM) - - job_keys_to_stop = job_keys or list(self.jobs.keys()) - for job_key in job_keys_to_stop: - for job_key_call_id in list(self.task_processes.keys()): - if job_key_call_id.rsplit('-', 1)[0] == job_key: - process = self.task_processes[job_key_call_id] - try: - kill_process(process) - except Exception: - pass - self.task_processes[job_key_call_id] = None - - super().stop(job_keys) + self._run_task_process( + f'{job_key}-{call_id}', + [self.runtime_name, RUNNER_FILE, 'run_job', task_filename], + ) class ContainerEnvironment(ExecutionEnvironment): - """ - Container environment uses a container runtime image - """ + """Container environment uses a container runtime image.""" - def __init__(self, config): + def __init__(self, config: Dict[str, Any]): super().__init__(config) logger.debug(f'Starting container environment for {self.runtime_name}') self.use_gpu = self.config.get('use_gpu', False) @@ -373,83 +457,90 @@ def __init__(self, config): self.uid = os.getuid() if self.is_unix_system else None self.gid = os.getgid() if self.is_unix_system else None + def _container_run_cmd(self, name, *, use_gpu=False, **kwargs) -> List[str]: + return docker_run_cmd( + self.docker_path, + self.runtime_name, + name=name, + tmp_path=Path(TEMP_DIR).as_posix(), + uid=self.uid, + gid=self.gid, + is_podman=self.is_podman, + use_gpu=use_gpu, + **kwargs, + ) + def setup(self): + """Installs the runner in the temp dir and pulls the image if asked""" logger.debug('Setting up container environment') self._copy_lithops_to_tmp() - if self.config.get('pull_runtime', False): logger.debug(f'Pulling runtime {self.runtime_name}') sp.run( - shlex.split(f'{self.docker_path} pull {self.runtime_name}'), - check=True, stdout=sp.PIPE, universal_newlines=True + docker_pull_cmd(self.docker_path, self.runtime_name), + check=True, + stdout=sp.PIPE, + universal_newlines=True, ) def get_metadata(self): - if not os.path.isfile(RUNNER_FILE): - self.setup() + """Asks the runtime image for the packages it has installed""" + self._ensure_runner() logger.debug(f"Extracting metadata from: {self.runtime_name}") - - tmp_path = Path(TEMP_DIR).as_posix() - - cmd = f'{self.docker_path} run --name lithops_metadata ' - cmd += f'--user {self.uid}:{self.gid} ' if self.is_unix_system and not self.is_podman else '' - cmd += f'--env USER={os.getenv("USER", "root")} ' - cmd += f'--rm -v {tmp_path}:/tmp --entrypoint "python3" ' - cmd += f'{self.runtime_name} /tmp/{USER_TEMP_DIR}/localhost-runner.py get_metadata' - process = sp.run( - shlex.split(cmd), check=True, stdout=sp.PIPE, - universal_newlines=True, start_new_session=True + self._container_run_cmd( + 'lithops_metadata', + container_args=[DOCKER_RUNNER_FILE, 'get_metadata'], + ), + check=True, + stdout=sp.PIPE, + universal_newlines=True, + start_new_session=True, ) - runtime_meta = json.loads(process.stdout.strip()) - - return runtime_meta + return json.loads(process.stdout.strip()) def start(self): - if not os.path.isfile(RUNNER_FILE): - self.setup() - - tmp_path = Path(TEMP_DIR).as_posix() - - cmd = f'{self.docker_path} run --name {self.container_name} ' - cmd += '--gpus all ' if self.use_gpu else '' - cmd += f'--user {self.uid}:{self.gid} ' if self.is_unix_system and not self.is_podman else '' - cmd += f'--env USER={os.getenv("USER", "root")} ' - cmd += f'--rm -v {tmp_path}:/tmp -it --detach ' - cmd += f'--entrypoint=/bin/bash {self.runtime_name}' - - self.container_process = sp.Popen(shlex.split(cmd), stdout=sp.DEVNULL, start_new_session=True) - self.container_process.communicate() # blocks until the process finishes - - super().start() - - def run_task(self, job_key, call_id): """ - Runs a task + Starts the container that will run every task of this executor, and + the consumer threads that feed it through docker exec """ - job_key_call_id = f'{job_key}-{call_id}' - docker_job_dir = f'/tmp/{USER_TEMP_DIR}/jobs/{job_key}' - docker_task_filename = f'{docker_job_dir}/{call_id}.task' + self._ensure_runner() - logger.debug(f"Going to execute task process {job_key_call_id}") - cmd = f'{self.docker_path} exec {self.container_name} /bin/bash -c ' - cmd += f'"python3 /tmp/{USER_TEMP_DIR}/localhost-runner.py ' - cmd += f'run_job {docker_task_filename}"' + self.container_process = sp.Popen( + self._container_run_cmd( + self.container_name, + extra_run_args=['-it', '--detach'], + entrypoint='/bin/bash', + use_gpu=self.use_gpu, + ), + stdout=sp.DEVNULL, + start_new_session=True, + ) + self.container_process.communicate() + super().start() - process = sp.Popen(shlex.split(cmd), stdout=sp.PIPE, stderr=sp.PIPE, start_new_session=True) - self.task_processes[job_key_call_id] = process - process.communicate() # blocks until the process finishes - if process.returncode != 0: - logger.error(f"Task process {job_key_call_id} failed with return code {process.returncode}") - logger.debug(f"Task process {job_key_call_id} finished") + def run_task(self, job_key: str, call_id: str) -> None: + """Runs one task inside the already running container""" + docker_task_filename = ( + f'/tmp/{USER_TEMP_DIR}/jobs/{job_key}/{call_id}.task' + ) + self._run_task_process( + f'{job_key}-{call_id}', + docker_exec_python_cmd( + self.docker_path, + self.container_name, + DOCKER_RUNNER_FILE, + 'run_job', + docker_task_filename, + ), + ) - def stop(self, job_keys=None): - """ - Stop localhost container - """ + def _teardown(self) -> None: + """Removes the container along with the consumer threads""" sp.Popen( - shlex.split(f'{self.docker_path} rm -f {self.container_name}'), - stdout=sp.DEVNULL, stderr=sp.DEVNULL + docker_rm_cmd(self.docker_path, self.container_name), + stdout=sp.DEVNULL, + stderr=sp.DEVNULL, ) - super().stop(job_keys) + super()._teardown() diff --git a/lithops/localhost/v2/runner.py b/lithops/localhost/v2/runner.py index a43bdde41..515e78a2f 100644 --- a/lithops/localhost/v2/runner.py +++ b/lithops/localhost/v2/runner.py @@ -20,79 +20,124 @@ import platform import logging import uuid +import traceback import multiprocessing as mp from lithops.worker import function_handler from lithops.worker.utils import get_runtime_metadata +from lithops.utils import log_prefix from lithops.constants import ( LITHOPS_TEMP_DIR, JOBS_DIR, LOGS_DIR, LOGGER_FORMAT, - RN_LOG_FILE + RN_LOG_FILE, ) - -os.makedirs(LITHOPS_TEMP_DIR, exist_ok=True) -os.makedirs(JOBS_DIR, exist_ok=True) -os.makedirs(LOGS_DIR, exist_ok=True) - -log_file_stream = open(RN_LOG_FILE, 'a') -logging.basicConfig(stream=log_file_stream, level=logging.DEBUG, format=LOGGER_FORMAT) logger = logging.getLogger('lithops.localhost.runner') -# Python 3.14 defaults to forkserver on Linux; Lithops requires fork. -if platform.system() != 'Windows': +def _configure_runner_logging(): + """ + Sends the runner logs to the runner log file, and returns the stream so + that the caller can also redirect the task output to it + """ + os.makedirs(LITHOPS_TEMP_DIR, exist_ok=True) + os.makedirs(JOBS_DIR, exist_ok=True) + os.makedirs(LOGS_DIR, exist_ok=True) + log_file_stream = open(RN_LOG_FILE, 'a') + logging.basicConfig( + stream=log_file_stream, + level=logging.DEBUG, + format=LOGGER_FORMAT, + ) + return log_file_stream + + +def _set_fork_start_method(): + """ + Forces fork, which Lithops relies on for its workers to inherit the task. + Python 3.14 defaults to forkserver on Linux + """ + if platform.system() == 'Windows': + return try: mp.set_start_method('fork') except RuntimeError: + # Already set by an earlier call in this interpreter pass -def run_job(): +def run_job(log_file_stream): + """ + Runs the single task described by the task file given as the second + argument + """ + # This process has no console: anything printed goes to the runner log sys.stdout = log_file_stream sys.stderr = log_file_stream task_filename = sys.argv[2] logger.info(f'Got {task_filename} file') - with open(task_filename, 'rb') as jf: - task_payload = json.load(jf) + with open(task_filename, 'r') as task_file: + task_payload = json.load(task_file) executor_id = task_payload['executor_id'] job_id = task_payload['job_id'] call_id = task_payload['call_ids'][0] - logger.info(f'ExecutorID {executor_id} | JobID {job_id} | CallID {call_id} - Starting execution') + logger.info( + f'{log_prefix(executor_id, job_id, call_id)} - Starting execution' + ) act_id = str(uuid.uuid4()).replace('-', '')[:12] os.environ['__LITHOPS_ACTIVATION_ID'] = act_id os.environ['__LITHOPS_BACKEND'] = 'Localhost' try: + # The environment already starts one runner per task, so the handler + # must not fork any further worker process task_payload['worker_processes'] = 1 function_handler(task_payload) except KeyboardInterrupt: pass - logger.info(f'ExecutorID {executor_id} | JobID {job_id} | CallID {call_id} - Execution Finished') + logger.info( + f'{log_prefix(executor_id, job_id, call_id)} - Execution Finished' + ) def extract_runtime_meta(): - runtime_meta = get_runtime_metadata() - print(json.dumps(runtime_meta)) + """Prints the metadata of this runtime, which the client reads back""" + print(json.dumps(get_runtime_metadata())) -if __name__ == "__main__": - logger.info('Starting Localhost task runner') - command = sys.argv[1] - logger.info(f'Received command: {command}') +def main(): + """Entry point of the runner subprocess, dispatching the argv command""" + _set_fork_start_method() + log_file_stream = _configure_runner_logging() + try: + logger.info('Starting Localhost task runner') + command = sys.argv[1] + logger.info(f'Received command: {command}') + + handlers = { + 'get_metadata': extract_runtime_meta, + 'run_job': lambda: run_job(log_file_stream), + } + handler = handlers.get(command) + if handler is None: + logger.error(f'Invalid command: {command}') + sys.exit(1) + handler() + except Exception: + logger.exception('Localhost task runner failed') + traceback.print_exc(file=sys.__stderr__) + sys.exit(1) + finally: + log_file_stream.close() - switcher = { - 'get_metadata': extract_runtime_meta, - 'run_job': run_job - } - switcher.get(command, lambda: "Invalid command")() - log_file_stream.close() +if __name__ == "__main__": + main() diff --git a/lithops/monitor.py b/lithops/monitor.py index f2c70d080..de6e50147 100644 --- a/lithops/monitor.py +++ b/lithops/monitor.py @@ -26,6 +26,8 @@ import concurrent.futures as cf from tblib import pickling_support +from lithops.utils import _future_id, log_prefix, monitoring_queue_name + pickling_support.install() logger = logging.getLogger(__name__) @@ -33,9 +35,26 @@ LOG_INTERVAL = 30 # Print monitor debug every LOG_INTERVAL seconds +def _status_id(call_status): + return ( + call_status['executor_id'], + call_status['job_id'], + call_status['call_id'], + ) + + +def _is_finished(fut): + return fut.ready or fut.success or fut.done + + +def _is_started(fut): + return fut.running or _is_finished(fut) + + class Monitor(threading.Thread): """ - Monitor base class + Base class of the background threads that follow the futures of an + executor and move them along their states as their status arrives """ def __init__(self, executor_id, @@ -66,44 +85,39 @@ def add_futures(self, fs): """ Extends the current thread list of futures to track """ - self.futures.update(set(fs)) - - present_jobs = {future.job_id for future in fs} - for job_id in present_jobs: - self.present_jobs.add(job_id) + self.futures.update(fs) + self.present_jobs.update(future.job_id for future in fs) def remove_futures(self, fs): """ Remove from the current thread a list of futures """ self._print_status_log() - - for future in fs: - if future in self.futures: - self.futures.remove(future) - - for job_id in {future.job_id for future in fs}: - if job_id in self.present_jobs: - self.present_jobs.remove(job_id) + self.futures.difference_update(fs) + self.present_jobs = {future.job_id for future in self.futures} def _all_ready(self): """ Checks if all futures are ready, success or done """ try: - return all(f.ready or f.success or f.done for f in self.futures) + return all(_is_finished(f) for f in self.futures) except Exception: + # Other threads add futures to the set while this one iterates + # it. A concurrent update means there is still work to wait for return False def _check_new_futures(self, call_status, f): - """Checks if a functions returned new futures to track""" + """ + Checks if a function returned new futures to track + """ if 'new_futures' not in call_status: return False f._set_futures(call_status) self.futures.update(f._new_futures) logger.debug( - f'ExecutorID {self.executor_id} - Received {len(f._new_futures)} ' + f'{log_prefix(self.executor_id)} - Received {len(f._new_futures)} ' 'new function Futures to track' ) @@ -120,37 +134,66 @@ def _future_timeout_checker(self, futures): start_tstamp = fut._call_status['worker_start_tstamp'] fut_timeout = start_tstamp + fut.execution_timeout + 5 if current_time > fut_timeout: - msg = f"The function exceeded the execution timeout of {fut.execution_timeout} seconds." + msg = ( + 'The function exceeded the execution timeout ' + f'of {fut.execution_timeout} seconds.' + ) raise TimeoutError('HANDLER', msg) except TimeoutError: - # generate fake TimeoutError call status + # Raising and catching the error right away is what fills + # sys.exc_info(), so that the client re-raises a real + # traceback for a worker that never reported back pickled_exception = str(pickle.dumps(sys.exc_info())) - call_status = {'type': '__end__', - 'exception': True, - 'exc_info': pickled_exception, - 'executor_id': fut.executor_id, - 'job_id': fut.job_id, - 'call_id': fut.call_id, - 'activation_id': fut.activation_id, - 'worker_start_tstamp': start_tstamp, - 'worker_end_tstamp': time.time()} + call_status = { + 'type': '__end__', + 'exception': True, + 'exc_info': pickled_exception, + 'executor_id': fut.executor_id, + 'job_id': fut.job_id, + 'call_id': fut.call_id, + 'activation_id': fut.activation_id, + 'worker_start_tstamp': start_tstamp, + 'worker_end_tstamp': time.time(), + } fut._set_ready(call_status) def _print_status_log(self, previous_log=None, log_time=None): - """prints a debug log showing the status of the job""" + """ + Logs how many calls are pending, running and done, but only when the + counts moved or the job has been silent for LOG_INTERVAL seconds + """ if not self.futures: return previous_log, log_time - callids_pending = len([f for f in self.futures if f.invoked]) - callids_running = len([f for f in self.futures if f.running]) - callids_done = len([f for f in self.futures if f.ready or f.success or f.done]) - if (callids_pending, callids_running, callids_done) != previous_log or log_time > LOG_INTERVAL: - logger.debug(f'ExecutorID {self.executor_id} - Pending: {callids_pending} ' - f'- Running: {callids_running} - Done: {callids_done}') + callids_pending = callids_running = callids_done = 0 + for fut in self.futures: + if fut.invoked: + callids_pending += 1 + if fut.running: + callids_running += 1 + if _is_finished(fut): + callids_done += 1 + counts = (callids_pending, callids_running, callids_done) + still_working = not all(_is_finished(fut) for fut in self.futures) + if counts != previous_log or ( + still_working + and log_time is not None + and log_time > LOG_INTERVAL + ): + logger.debug( + f'{log_prefix(self.executor_id)} - Pending: ' + f'{callids_pending} - Running: {callids_running} - Done: {callids_done}' + ) log_time = 0 - return (callids_pending, callids_running, callids_done), log_time + return counts, log_time class RabbitmqMonitor(Monitor): + """ + Job monitor that learns the status of every call from the messages the + workers publish to a RabbitMQ queue + """ + + SLEEP_TIME = 2 def __init__( self, @@ -171,7 +214,7 @@ def __init__( ) self.rabbit_amqp_url = config.get('amqp_url') - self.queue = f'lithops-{self.executor_id}' + self.queue = monitoring_queue_name(self.executor_id) self.tag = None self._create_resources() @@ -179,7 +222,9 @@ def _create_resources(self): """ Creates RabbitMQ queues and exchanges of a given job """ - logger.debug(f'ExecutorID {self.executor_id} - Creating RabbitMQ queue {self.queue}') + logger.debug( + f'{log_prefix(self.executor_id)} - Creating RabbitMQ queue {self.queue}' + ) self.pikaparams = pika.URLParameters(self.rabbit_amqp_url) self.connection = pika.BlockingConnection(self.pikaparams) @@ -210,82 +255,105 @@ def _tag_future_as_running(self, call_status): """ Assigns a call_status to its future """ - not_running_futures = [f for f in self.futures if not (f.running or f.ready or f.success or f.done)] + not_running_futures = [ + f for f in self.futures if not _is_started(f) + ] for f in not_running_futures: - calljob_id = (call_status['executor_id'], call_status['job_id'], call_status['call_id']) - if (f.executor_id, f.job_id, f.call_id) == calljob_id: + if _future_id(f) == _status_id(call_status): f._set_running(call_status) def _tag_future_as_ready(self, call_status): """ - tags a future as ready based on call_status + Tags a future as ready based on call_status """ - not_ready_futures = [f for f in self.futures if not (f.ready or f.success or f.done)] + not_ready_futures = [ + f for f in self.futures if not _is_finished(f) + ] for f in not_ready_futures: - calljob_id = (call_status['executor_id'], call_status['job_id'], call_status['call_id']) - if (f.executor_id, f.job_id, f.call_id) == calljob_id: + if _future_id(f) == _status_id(call_status): if not self._check_new_futures(call_status, f): f._set_ready(call_status) def _generate_tokens(self, call_status): """ - generates a new token for the invoker + Hands a token back to the invoker once a whole worker is free """ if not self.generate_tokens or not self.should_run: return - call_id = (call_status['executor_id'], call_status['job_id'], call_status['call_id']) + call_id = _status_id(call_status) worker_id = call_status['activation_id'] - if worker_id not in self.callids_done_worker: - self.callids_done_worker[worker_id] = [] - self.callids_done_worker[worker_id].append(call_id) + done_for_worker = self.callids_done_worker.setdefault(worker_id, []) + done_for_worker.append(call_id) - if worker_id not in self.workers_done and \ - len(self.callids_done_worker[worker_id]) == call_status['chunksize']: + if ( + worker_id not in self.workers_done + and len(done_for_worker) == call_status['chunksize'] + ): self.workers_done.append(worker_id) if self.should_run: self.token_bucket_q.put('#') - def run(self): - logger.debug(f'ExecutorID {self.executor_id} | Starting RabbitMQ job monitor') - SLEEP_TIME = 2 - - channel = self.connection.channel() + def _on_message(self, ch, method, properties, body): + """ + Applies one status message to its future, and stops consuming once + there is nothing left to wait for + """ + call_status = json.loads(body.decode("utf-8")) - def callback(ch, method, properties, body): - call_status = json.loads(body.decode("utf-8")) + if call_status['type'] == '__init__': + self._tag_future_as_running(call_status) - if call_status['type'] == '__init__': - self._tag_future_as_running(call_status) + elif call_status['type'] == '__end__': + self._generate_tokens(call_status) + self._tag_future_as_ready(call_status) - elif call_status['type'] == '__end__': - self._generate_tokens(call_status) - self._tag_future_as_ready(call_status) + if self._all_ready() or not self.should_run: + ch.stop_consuming() + ch.close() - if self._all_ready() or not self.should_run: - ch.stop_consuming() - ch.close() + def _watch_timeouts(self): + """ + Logs the job status and expires overdue futures. Runs in its own + thread, as the monitor thread stays blocked on the queue + """ + previous_log = None + log_time = 0 + while self.should_run and not self._all_ready(): + previous_log, log_time = self._print_status_log( + previous_log=previous_log, log_time=log_time + ) + self._future_timeout_checker(self.futures) + time.sleep(self.SLEEP_TIME) + log_time += self.SLEEP_TIME - def manage_timeouts(): - prevoius_log = None - log_time = 0 - while self.should_run and not self._all_ready(): - # Format call_ids running, pending and done - prevoius_log, log_time = self._print_status_log(previous_log=prevoius_log, log_time=log_time) - self._future_timeout_checker(self.futures) - time.sleep(SLEEP_TIME) - log_time += SLEEP_TIME + def run(self): + """ + Consumes status messages from the queue until every future is done + """ + logger.debug( + f'{log_prefix(self.executor_id)} | Starting RabbitMQ job monitor' + ) - threading.Thread(target=manage_timeouts, daemon=True).start() + channel = self.connection.channel() + threading.Thread(target=self._watch_timeouts, daemon=True).start() - self.tag = channel.basic_consume(self.queue, callback, auto_ack=True) + self.tag = channel.basic_consume( + self.queue, self._on_message, auto_ack=True + ) channel.start_consuming() self.tag = None self._print_status_log() - logger.debug(f'ExecutorID {self.executor_id} | RabbitMQ job monitor finished') + logger.debug( + f'{log_prefix(self.executor_id)} | RabbitMQ job monitor finished' + ) class StorageMonitor(Monitor): + """ + Job monitor that learns the status of every call by polling the storage + backend, where the workers leave their status objects + """ THREADPOOL_SIZE = 64 @@ -319,6 +387,7 @@ def __init__( # vars for _mark_status_as_ready self.callids_done_processed_status = set() + self._ready_pool = None def stop(self): """ @@ -326,30 +395,63 @@ def stop(self): """ self.should_run = False + def join(self, timeout=None): + """ + Waits for the monitor thread, and drops the pool it downloads the + statuses with, which outlives the thread on a join that timed out + """ + super().join(timeout) + self._shutdown_ready_pool() + + def _get_ready_pool(self): + if self._ready_pool is None: + self._ready_pool = cf.ThreadPoolExecutor( + max_workers=self.THREADPOOL_SIZE + ) + return self._ready_pool + + def _shutdown_ready_pool(self): + pool = self._ready_pool + if pool is None: + return + self._ready_pool = None + pool.shutdown(wait=False) + def _tag_future_as_running(self, callids_running): """ Mark which futures are in running status based on callids_running """ current_time = time.time() - not_running_futures = [f for f in self.futures if not (f.running or f.ready or f.success or f.done)] - callids_running_to_process = callids_running - self.callids_running_processed_timeout - for f in not_running_futures: - for call in callids_running_to_process: - if f.invoked and (f.executor_id, f.job_id, f.call_id) == call[0]: - call_status = {'type': '__init__', - 'activation_id': call[1], - 'worker_start_tstamp': current_time} - f._set_running(call_status) - - self.callids_running_processed_timeout.update(callids_running_to_process) + to_process = ( + callids_running - self.callids_running_processed_timeout + ) + pending = { + _future_id(f): f + for f in self.futures + if f.invoked and not _is_started(f) + } + for call in to_process: + f = pending.get(call[0]) + if f is None: + continue + call_status = { + 'type': '__init__', + 'activation_id': call[1], + 'worker_start_tstamp': current_time, + } + f._set_running(call_status) + + self.callids_running_processed_timeout.update(to_process) self._future_timeout_checker(self.futures) def _tag_future_as_ready(self, callids_done): """ Mark which futures has a call_status ready to be downloaded """ - not_ready_futures = [f for f in self.futures if not (f.ready or f.success or f.done)] - callids_done_to_process = callids_done - self.callids_done_processed_status + not_ready_futures = [ + f for f in self.futures if not _is_finished(f) + ] + to_process = callids_done - self.callids_done_processed_status fs_to_query = [] ten_percent = int(len(self.futures) * (10 / 100)) @@ -357,90 +459,91 @@ def _tag_future_as_ready(self, callids_done): fs_to_query = not_ready_futures else: for f in not_ready_futures: - if (f.executor_id, f.job_id, f.call_id) in callids_done_to_process: + if _future_id(f) in to_process: fs_to_query.append(f) if not fs_to_query: return def get_status(f): - cs = self.internal_storage.get_call_status(f.executor_id, f.job_id, f.call_id) + cs = self.internal_storage.get_call_status( + f.executor_id, f.job_id, f.call_id + ) f._status_query_count += 1 if cs: if not self._check_new_futures(cs, f): f._set_ready(cs) - return (f.executor_id, f.job_id, f.call_id) - else: - return None + return _future_id(f) + return None try: - pool = cf.ThreadPoolExecutor(max_workers=self.THREADPOOL_SIZE) - call_ids_processed = set(pool.map(get_status, fs_to_query)) - pool.shutdown() - except Exception: - pass - - try: - call_ids_processed.remove(None) + call_ids_processed = set( + self._get_ready_pool().map(get_status, fs_to_query) + ) except Exception: - pass + return + finally: + # The final sweep of run() happens after the thread is done, so + # the pool it lazily recreated has to be dropped again + if not self.is_alive(): + self._shutdown_ready_pool() - try: - self.callids_done_processed_status.update(call_ids_processed) - except Exception: - pass + call_ids_processed.discard(None) + self.callids_done_processed_status.update(call_ids_processed) def _generate_tokens(self, callids_running, callids_done): """ - Method that generates new tokens + Hands a token back to the invoker for every worker that finished the + whole chunk of calls it was given """ if not self.generate_tokens or not self.should_run: return - callids_running_to_process = callids_running - self.callids_running_processed - callids_done_to_process = callids_done - self.callids_done_processed + running_new = ( + callids_running - self.callids_running_processed + ) + done_new = callids_done - self.callids_done_processed - for call_id, worker_id in callids_running_to_process: - if worker_id not in self.workers: - self.workers[worker_id] = set() - self.workers[worker_id].add(call_id) + for call_id, worker_id in running_new: + self.workers.setdefault(worker_id, set()).add(call_id) self.callids_running_worker[call_id] = worker_id - for callid_done in callids_done_to_process: + for callid_done in done_new: if callid_done in self.callids_running_worker: worker_id = self.callids_running_worker[callid_done] - if worker_id not in self.callids_done_worker: - self.callids_done_worker[worker_id] = [] - self.callids_done_worker[worker_id].append(callid_done) + self.callids_done_worker.setdefault(worker_id, []).append( + callid_done + ) for worker_id in self.callids_done_worker: job_id = self.callids_done_worker[worker_id][0][1] if job_id not in self.present_jobs: continue chunksize = self.job_chunksize[job_id] - if worker_id not in self.workers_done and \ - len(self.callids_done_worker[worker_id]) == chunksize: + done_count = len(self.callids_done_worker[worker_id]) + if worker_id not in self.workers_done and done_count == chunksize: self.workers_done.append(worker_id) if self.should_run: self.token_bucket_q.put('#') else: break - self.callids_running_processed.update(callids_running_to_process) - self.callids_done_processed.update(callids_done_to_process) + self.callids_running_processed.update(running_new) + self.callids_done_processed.update(done_new) def _poll_and_process_job_status(self, previous_log, log_time): """ - Polls the storage backend for job status, updates futures, - and prints status logs. - - Returns: - new_callids_done (set): New callids that were marked as done. - previous_log (str): Updated log message. - log_time (float): Updated log time counter. + Reads the job status from storage and applies it to the futures. + Returns the call ids that are newly done, along with the updated + log state its caller has to pass back on the next round """ - callids_running, callids_done = self.internal_storage.get_job_status(self.executor_id) - new_callids_done = callids_done - self.callids_done_processed_status + status = self.internal_storage.get_job_status( + self.executor_id, job_ids=self.present_jobs + ) + callids_running, callids_done = status + new_callids_done = ( + callids_done - self.callids_done_processed_status + ) self._generate_tokens(callids_running, callids_done) self._tag_future_as_running(callids_running) @@ -452,9 +555,12 @@ def _poll_and_process_job_status(self, previous_log, log_time): def run(self): """ - Run method for the Storage job monitor thread. + Polls the storage backend until the monitor is stopped, backing off + to the configured interval whenever a round brings nothing new """ - logger.debug(f'ExecutorID {self.executor_id} - Starting Storage job monitor') + logger.debug( + f'{log_prefix(self.executor_id)} - Starting Storage job monitor' + ) wait_dur_sec = self.monitoring_interval previous_log = None @@ -462,22 +568,44 @@ def run(self): while self.should_run: try: - new_callids_done, previous_log, log_time = self._poll_and_process_job_status(previous_log, log_time) + new_callids_done, previous_log, log_time = ( + self._poll_and_process_job_status( + previous_log, log_time + ) + ) if new_callids_done: wait_dur_sec = self.monitoring_interval / 5 else: wait_dur_sec = self.monitoring_interval except Exception as e: - logger.error(f'ExecutorID {self.executor_id} - Error during monitor: {e}', exc_info=True) + logger.error( + f'{log_prefix(self.executor_id)} - Error during ' + f'monitor: {e}', + exc_info=True, + ) + if not self.should_run: + break time.sleep(wait_dur_sec) log_time += wait_dur_sec - self._poll_and_process_job_status(previous_log, log_time) + # One last sweep, so that statuses written between the final poll + # and the stop are not lost. The storage may already be gone + try: + self._poll_and_process_job_status(previous_log, log_time) + except Exception: + pass - logger.debug(f'ExecutorID {self.executor_id} - Storage job monitor finished') + self._shutdown_ready_pool() + logger.debug( + f'{log_prefix(self.executor_id)} - Storage job monitor finished' + ) class JobMonitor: + """ + Owns the monitor thread of one executor, and picks the implementation + that matches the configured monitoring backend + """ def __init__(self, executor_id, internal_storage, config=None): self.executor_id = executor_id @@ -485,7 +613,9 @@ def __init__(self, executor_id, internal_storage, config=None): self.storage_config = internal_storage.get_storage_config() self.storage_backend = internal_storage.backend self.config = config - self.type = self.config['lithops']['monitoring'].lower() if config else 'storage' + self.type = ( + config['lithops']['monitoring'].lower() if config else 'storage' + ) self.token_bucket_q = queue.Queue() self.monitor = None @@ -497,9 +627,13 @@ def __init__(self, executor_id, internal_storage, config=None): ) def start(self, fs, job_id=None, chunksize=None, generate_tokens=False): + """ + Tracks a new set of futures, spawning the monitor thread unless a + live one can take them over + """ if self.type == 'storage': - monitoring_interval = self.storage_config['monitoring_interval'] - monitor_config = {'monitoring_interval': monitoring_interval} + interval = self.storage_config['monitoring_interval'] + monitor_config = {'monitoring_interval': interval} else: monitor_config = self.config.get(self.type) @@ -522,12 +656,22 @@ def start(self, fs, job_id=None, chunksize=None, generate_tokens=False): self.monitor.start() def is_alive(self): + """ + Tells whether the monitor thread is still running + """ return self.monitor.is_alive() def remove(self, fs): + """ + Stops tracking a set of futures + """ if self.monitor and self.monitor.is_alive(): self.monitor.remove_futures(fs) def stop(self): + """ + Stops the monitor thread and waits for it to wind down + """ if self.monitor and self.monitor.is_alive(): self.monitor.stop() + self.monitor.join(timeout=5) diff --git a/lithops/plots.py b/lithops/plots.py index 5cdf654ca..25a16cef6 100644 --- a/lithops/plots.py +++ b/lithops/plots.py @@ -30,9 +30,81 @@ logger = logging.getLogger(__name__) +def _plot_destination(dst, suffix): + """ + Resolves where a plot has to be written, defaulting to a timestamped + file under a plots directory of the working directory + """ + if dst is None: + os.makedirs('plots', exist_ok=True) + filename = f'{int(time.time())}_{suffix}' + return os.path.join(os.getcwd(), 'plots', filename) + return f'{os.path.realpath(os.path.expanduser(dst))}_{suffix}' + + +def _set_call_axis(ax, total_calls): + yplot_step = max(1, total_calls // 20) + y_ticks = np.arange(total_calls // yplot_step + 2) * yplot_step + ax.set_yticks(y_ticks) + ax.set_ylim(-0.02 * total_calls, total_calls * 1.02) + return y_ticks + + +def _set_time_axis(ax, max_seconds): + xplot_step = max(int(max_seconds / 8), 1) + x_ticks = np.arange(max_seconds // xplot_step + 2) * xplot_step + ax.set_xlim(0, max_seconds) + ax.set_xticks(x_ticks) + for x in x_ticks: + ax.axvline(x, c='k', alpha=0.2, linewidth=0.8) + return x_ticks + + +def _elapsed(series, t0): + return series - t0 + + +def _timeline_span(stats_df, t0): + """ + Returns how far the time axis has to reach, based on the last milestone + the run got to record + """ + if 'host_result_done_tstamp' in stats_df: + col = stats_df.host_result_done_tstamp + elif 'host_status_done_tstamp' in stats_df: + col = stats_df.host_status_done_tstamp + else: + col = stats_df.end_tstamp + return np.max(col - t0) * 1.25 + + +def _timeline_fields(stats_df, t0): + """ + Returns a (label, elapsed times) pair per milestone that every call goes + through. Results are only there when the client downloaded them + """ + fields = [ + ('host submit', _elapsed(stats_df.host_submit_tstamp, t0)), + ('function start', _elapsed(stats_df.worker_func_start_tstamp, t0)), + ('function done', _elapsed(stats_df.worker_func_end_tstamp, t0)), + ('status fetched', _elapsed(stats_df.host_status_done_tstamp, t0)), + ] + + if 'host_result_done_tstamp' in stats_df: + fields.append( + ('results fetched', _elapsed(stats_df.host_result_done_tstamp, t0)) + ) + + return fields + + def create_timeline(fs, dst, figsize=(10, 6)): + """ + Plots when every call reached each of its milestones, and writes the + figure next to dst + """ stats = [f.stats for f in fs] - host_job_create_tstamp = min([cm['host_job_create_tstamp'] for cm in stats]) + t0 = min(cm['host_job_create_tstamp'] for cm in stats) stats_df = pd.DataFrame(stats) total_calls = len(stats_df) @@ -44,133 +116,108 @@ def create_timeline(fs, dst, figsize=(10, 6)): y = np.arange(total_calls) point_size = 10 - - fields = [('host submit', stats_df.host_submit_tstamp - host_job_create_tstamp), - # ('worker start', stats_df.worker_start_tstamp - host_job_create_tstamp), - ('function start', stats_df.worker_func_start_tstamp - host_job_create_tstamp), - ('function done', stats_df.worker_func_end_tstamp - host_job_create_tstamp), - # ('worker done', stats_df.worker_end_tstamp - host_job_create_tstamp), - ('status fetched', stats_df.host_status_done_tstamp - host_job_create_tstamp)] - - if 'host_result_done_tstamp' in stats_df: - fields.append(('results fetched', stats_df.host_result_done_tstamp - host_job_create_tstamp)) + fields = _timeline_fields(stats_df, t0) patches = [] for f_i, (field_name, val) in enumerate(fields): - ax.scatter(val, y, c=[palette[f_i]], edgecolor='none', s=point_size, alpha=0.8) - patches.append(mpatches.Patch(color=palette[f_i], label=field_name)) + ax.scatter( + val, y, c=[palette[f_i]], + edgecolor='none', s=point_size, alpha=0.8, + ) + patches.append( + mpatches.Patch(color=palette[f_i], label=field_name) + ) ax.set_xlabel('Execution Time (sec)') ax.set_ylabel('Function Call') - legend = pylab.legend(handles=patches, loc='upper right', frameon=True) + legend = pylab.legend( + handles=patches, loc='upper right', frameon=True + ) legend.get_frame().set_facecolor('#FFFFFF') - yplot_step = int(np.max([1, total_calls / 20])) - y_ticks = np.arange(total_calls // yplot_step + 2) * yplot_step - ax.set_yticks(y_ticks) - ax.set_ylim(-0.02 * total_calls, total_calls * 1.02) - for y in y_ticks: - ax.axhline(y, c='k', alpha=0.1, linewidth=1) + y_ticks = _set_call_axis(ax, total_calls) + for ytick in y_ticks: + ax.axhline(ytick, c='k', alpha=0.1, linewidth=1) - if 'host_result_done_tstamp' in stats_df: - max_seconds = np.max(stats_df.host_result_done_tstamp - host_job_create_tstamp) * 1.25 - elif 'host_status_done_tstamp' in stats_df: - max_seconds = np.max(stats_df.host_status_done_tstamp - host_job_create_tstamp) * 1.25 - else: - max_seconds = np.max(stats_df.end_tstamp - host_job_create_tstamp) * 1.25 - xplot_step = max(int(max_seconds / 8), 1) - x_ticks = np.arange(max_seconds // xplot_step + 2) * xplot_step - ax.set_xlim(0, max_seconds) - - ax.set_xticks(x_ticks) - for x in x_ticks: - ax.axvline(x, c='k', alpha=0.2, linewidth=0.8) + _set_time_axis(ax, _timeline_span(stats_df, t0)) ax.grid(False) fig.tight_layout() + fig.savefig(_plot_destination(dst, 'timeline.png')) - if dst is None: - os.makedirs('plots', exist_ok=True) - dst = os.path.join(os.getcwd(), 'plots', '{}_{}'.format(int(time.time()), 'timeline.png')) - else: - dst = os.path.expanduser(dst) if '~' in dst else dst - dst = '{}_{}'.format(os.path.realpath(dst), 'timeline.png') - fig.savefig(dst) +def _active_calls_hist(time_rates, t0, runtime_bins): + """ + Turns (start, end) pairs into elapsed times, plus a per call mask of the + time bins the call was active in, which summed gives the concurrency + """ + x = np.array(time_rates) + start_time = x[:, 0] - t0 + end_time = x[:, 1] - t0 + + calls_hist = np.zeros((len(start_time), len(runtime_bins))) + for i, (start, end) in enumerate(zip(start_time, end_time)): + a, b = np.searchsorted(runtime_bins, [start, end]) + if b - a > 0: + calls_hist[i, a:b] = 1 + + return start_time, end_time, calls_hist def create_histogram(fs, dst, figsize=(10, 6)): + """ + Plots how long every call ran for, over the number of calls that were + running at the same time, and writes the figure next to dst + """ stats = [f.stats for f in fs] - host_job_create_tstamp = min([cm['host_job_create_tstamp'] for cm in stats]) + t0 = min(cm['host_job_create_tstamp'] for cm in stats) total_calls = len(stats) - max_seconds = int(max([cs['worker_end_tstamp'] - host_job_create_tstamp for cs in stats]) * 2.5) + max_seconds = int(max( + cs['worker_end_tstamp'] - t0 for cs in stats + ) * 2.5) runtime_bins = np.linspace(0, max_seconds, max_seconds) - def compute_times_rates(time_rates): - x = np.array(time_rates) - tzero = host_job_create_tstamp - start_time = x[:, 0] - tzero - end_time = x[:, 1] - tzero - - N = len(start_time) - - runtime_calls_hist = np.zeros((N, len(runtime_bins))) - - for i in range(N): - s = start_time[i] - e = end_time[i] - a, b = np.searchsorted(runtime_bins, [s, e]) - if b - a > 0: - runtime_calls_hist[i, a:b] = 1 - - return {'start_tstamp': start_time, - 'end_tstamp': end_time, - 'runtime_calls_hist': runtime_calls_hist} - fig = pylab.figure(figsize=figsize) ax = fig.add_subplot(1, 1, 1) - time_rates = [(cs['worker_start_tstamp'], cs['worker_end_tstamp']) for cs in stats] - - time_hist = compute_times_rates(time_rates) - - N = len(time_hist['start_tstamp']) - line_segments = LineCollection([[[time_hist['start_tstamp'][i], i], - [time_hist['end_tstamp'][i], i]] for i in range(N)], - linestyles='solid', color='k', alpha=0.6, linewidth=0.4) - - ax.add_collection(line_segments) - - ax.plot(runtime_bins, time_hist['runtime_calls_hist'].sum(axis=0), label='Total Active Calls', zorder=-1) + time_rates = [ + (cs['worker_start_tstamp'], cs['worker_end_tstamp']) + for cs in stats + ] + start_time, end_time, calls_hist = _active_calls_hist( + time_rates, t0, runtime_bins + ) + + segments = [ + [[start_time[i], i], [end_time[i], i]] + for i in range(len(start_time)) + ] + ax.add_collection(LineCollection( + segments, + linestyles='solid', + color='k', + alpha=0.6, + linewidth=0.4, + )) + + ax.plot( + runtime_bins, + calls_hist.sum(axis=0), + label='Total Active Calls', + zorder=-1, + ) + + _set_call_axis(ax, total_calls) + _set_time_axis(ax, max_seconds) - yplot_step = int(np.max([1, total_calls / 20])) - y_ticks = np.arange(total_calls // yplot_step + 2) * yplot_step - ax.set_yticks(y_ticks) - ax.set_ylim(-0.02 * total_calls, total_calls * 1.02) - - xplot_step = max(int(max_seconds / 8), 1) - x_ticks = np.arange(max_seconds // xplot_step + 2) * xplot_step - ax.set_xlim(0, max_seconds) - ax.set_xticks(x_ticks) - for x in x_ticks: - ax.axvline(x, c='k', alpha=0.2, linewidth=0.8) - - ax.set_xlabel("Execution Time (sec)") - ax.set_ylabel("Function Call") + ax.set_xlabel('Execution Time (sec)') + ax.set_ylabel('Function Call') ax.grid(False) ax.legend(loc='upper right') fig.tight_layout() - - if dst is None: - os.makedirs('plots', exist_ok=True) - dst = os.path.join(os.getcwd(), 'plots', '{}_{}'.format(int(time.time()), 'histogram.png')) - else: - dst = os.path.expanduser(dst) if '~' in dst else dst - dst = '{}_{}'.format(os.path.realpath(dst), 'histogram.png') - - fig.savefig(dst) + fig.savefig(_plot_destination(dst, 'histogram.png')) pylab.close(fig) diff --git a/lithops/retries.py b/lithops/retries.py index 57e7e480e..395b84a4e 100644 --- a/lithops/retries.py +++ b/lithops/retries.py @@ -62,73 +62,61 @@ def __init__( self.cancelled = False def _inc_failure_count(self): - """ - Increment the internal failure counter. - """ self.failure_count += 1 def _should_retry(self): """ - Determine whether another retry attempt should be made. - - :return: True if retry is allowed, False otherwise. + Tells whether the call has retries left and was not cancelled """ return not self.cancelled and self.failure_count <= self.retries def _retry(self, function_executor: FunctionExecutor): """ - Re-submit the map function using the provided FunctionExecutor. - - :param function_executor: An instance of FunctionExecutor to resubmit the job. + Resubmits the map function with the same input, and takes the new + activation as the future to follow from now on """ - inputs = [self.input] - futures_list = function_executor.map( - self.map_function, inputs, **self.map_kwargs - ) - self.response_future = futures_list[0] + self.response_future = function_executor.map( + self.map_function, [self.input], **self.map_kwargs + )[0] def cancel(self): """ - Cancel any future retries. This does not cancel any running tasks. + Gives up on any further retry. A task already running is left alone """ self.cancelled = True @property def done(self): """ - Check if the function execution is complete. - - :return: True if the execution is done, False otherwise. + Tells whether the function execution is complete """ return self.response_future.done @property def error(self): """ - Get the error from the function execution, if any. - - :return: An exception or error message if an error occurred. + Tells whether the function execution raised """ return self.response_future.error @property def _exception(self): """ - Get the exception tuple (type, value, traceback) from the execution. - - :return: Exception tuple. + Returns the (type, value, traceback) tuple the execution raised """ return self.response_future._exception @property def stats(self): """ - Get execution statistics. - - :return: A dictionary containing performance and usage metrics. + Returns the performance and usage metrics of the execution """ return self.response_future.stats + def _reraise_if_error(self): + if self.response_future.error: + reraise(*self.response_future._exception) + def status( self, throw_except: bool = True, @@ -148,8 +136,7 @@ def status( internal_storage=internal_storage, check_only=check_only, ) - if self.response_future.error: - reraise(*self.response_future._exception) + self._reraise_if_error() return stat def result(self, throw_except: bool = True, internal_storage: Any = None): @@ -163,21 +150,19 @@ def result(self, throw_except: bool = True, internal_storage: Any = None): res = self.response_future.result( throw_except=throw_except, internal_storage=internal_storage ) - if self.response_future.error: - reraise(*self.response_future._exception) + self._reraise_if_error() return res class RetryingFunctionExecutor: """ - A wrapper around `FunctionExecutor` that adds automatic retry capabilities to function invocations. - This class allows spawning multiple function activations and handling failures by retrying them - according to the configured retry policy. + Wrapper around FunctionExecutor with automatic retries. - It provides the same interface as `FunctionExecutor` for compatibility, with an extra `retries` parameter - in `map()` to control the number of retries per invocation. + Same public interface as FunctionExecutor, plus a `retries` + parameter on map() for per-invocation retry budget. - :param executor: An instance of FunctionExecutor (e.g., Localhost, Serverless, or Standalone) + :param executor: FunctionExecutor (localhost, serverless, + or standalone) """ def __init__(self, executor: FunctionExecutor): @@ -197,12 +182,21 @@ def __exit__(self, exc_type, exc_value, traceback): """ self.executor.__exit__(exc_type, exc_value, traceback) + def _retries_to_use(self, retries): + if retries is not None: + return retries + return self.config.get('lithops', {}).get('retries', 0) + def map( self, map_function: Callable, - map_iterdata: List[Union[List[Any], Tuple[Any, ...], Dict[str, Any]]], + map_iterdata: List[Union[ + List[Any], Tuple[Any, ...], Dict[str, Any] + ]], chunksize: Optional[int] = None, - extra_args: Optional[Union[List[Any], Tuple[Any, ...], Dict[str, Any]]] = None, + extra_args: Optional[Union[ + List[Any], Tuple[Any, ...], Dict[str, Any] + ]] = None, extra_env: Optional[Dict[str, str]] = None, runtime_memory: Optional[int] = None, obj_chunk_size: Optional[int] = None, @@ -218,30 +212,36 @@ def map( :param map_function: The function to map over the data. :param map_iterdata: An iterable of input data (e.g., Python list). - :param chunksize: Split map_iterdata in chunks of this size. One worker per chunk. - :param extra_args: Additional arguments to pass to each function. - :param extra_env: Additional environment variables for the function environment. - :param runtime_memory: Memory (in MB) to allocate per function activation. - :param obj_chunk_size: For file processing. Split each object into chunks of this size (in bytes). - :param obj_chunk_number: For file processing. Number of chunks to split each object into. - :param obj_newline: Newline character for line integrity in file partitioning. - :param timeout: Max time per function activation (in seconds). - :param include_modules: Explicitly pickle these dependencies. - :param exclude_modules: Explicitly exclude these modules from pickling. - :param retries: Number of retries for each function activation upon failure. - - :return: A list of RetryingFuture objects, one for each function activation. - """ - - retries_to_use = ( - retries - if retries is not None - else self.config.get('lithops', {}).get('retries', 0) - ) - - futures_list = self.executor.map( - map_function, - map_iterdata, + :param chunksize: Split map_iterdata in chunks of this + size. One worker per chunk. + :param extra_args: Additional arguments to pass to each + function. + :param extra_env: Additional environment variables for + the function environment. + :param runtime_memory: Memory (in MB) to allocate per + function activation. + :param obj_chunk_size: For file processing. Split each + object into chunks of this size (in bytes). + :param obj_chunk_number: For file processing. Number of + chunks to split each object into. + :param obj_newline: Newline character for line integrity + in file partitioning. + :param timeout: Max time per function activation + (in seconds). + :param include_modules: Explicitly pickle these + dependencies. + :param exclude_modules: Explicitly exclude these modules + from pickling. + :param retries: Number of retries for each function + activation upon failure. + + :return: A list of RetryingFuture objects, one for each + function activation. + """ + + retries_to_use = self._retries_to_use(retries) + + map_kwargs = dict( chunksize=chunksize, extra_args=extra_args, extra_env=extra_env, @@ -253,26 +253,65 @@ def map( include_modules=include_modules, exclude_modules=exclude_modules, ) + + futures_list = self.executor.map( + map_function, + map_iterdata, + **map_kwargs, + ) return [ RetryingFuture( f, map_function=map_function, input=i, retries=retries_to_use, - chunksize=chunksize, - extra_args=extra_args, - extra_env=extra_env, - runtime_memory=runtime_memory, - obj_chunk_size=obj_chunk_size, - obj_chunk_number=obj_chunk_number, - obj_newline=obj_newline, - timeout=timeout, - include_modules=include_modules, - exclude_modules=exclude_modules, + **map_kwargs, ) for i, f in zip(map_iterdata, futures_list) ] + def _split_done_and_retried(self, done, pending, lookup): + """ + Sorts the futures the inner executor reported as done into the ones + that are really finished and the ones that failed and got + resubmitted, extending the lookup with every new response future + """ + retrying_done = [] + retrying_pending = [ + lookup[response_future] for response_future in pending + ] + + for response_future in done: + retrying_future = lookup[response_future] + if not response_future.error: + retrying_done.append(retrying_future) + continue + + retrying_future._inc_failure_count() + if not retrying_future._should_retry(): + retrying_done.append(retrying_future) + continue + + retrying_future._retry(self.executor) + retrying_pending.append(retrying_future) + lookup[retrying_future.response_future] = retrying_future + + return retrying_done, retrying_pending + + @staticmethod + def _wait_is_over(return_when, retrying_done, retrying_pending): + """ + Tells whether the completion policy is satisfied, which for anything + but ALWAYS may take several rounds because of the retries + """ + if return_when == ALWAYS: + return True + if return_when == ANY_COMPLETED: + return bool(retrying_done) + if return_when == ALL_COMPLETED: + return not retrying_pending + return False + def wait( self, fs: List[RetryingFuture], @@ -289,7 +328,8 @@ def wait( :param fs: List of RetryingFuture objects to wait on. :param throw_except: Raise exceptions encountered during execution. - :param return_when: Completion policy. One of: ALWAYS, ANY_COMPLETED, or ALL_COMPLETED. + :param return_when: Completion policy. One of: ALWAYS, + ANY_COMPLETED, or ALL_COMPLETED. :param download_results: Whether to download results after completion. :param timeout: Maximum wait time (in seconds). :param threadpool_size: Number of threads used for polling. @@ -301,10 +341,10 @@ def wait( lookup = {f.response_future: f for f in fs} while True: - response_futures = [f.response_future for f in fs] - + # A retry replaces the response future of a RetryingFuture, so + # the list has to be rebuilt on every round done, pending = self.executor.wait( - response_futures, + [f.response_future for f in fs], throw_except=throw_except, return_when=return_when, download_results=download_results, @@ -314,29 +354,14 @@ def wait( show_progressbar=show_progressbar, ) - retrying_done = [] - retrying_pending = [lookup[response_future] for response_future in pending] - for response_future in done: - retrying_future = lookup[response_future] - if response_future.error: - retrying_future._inc_failure_count() - if retrying_future._should_retry(): - retrying_future._retry(self.executor) - retrying_pending.append(retrying_future) - lookup[retrying_future.response_future] = retrying_future - else: - retrying_done.append(retrying_future) - else: - retrying_done.append(retrying_future) - - if return_when == ALWAYS: - break - elif return_when == ANY_COMPLETED and len(retrying_done) > 0: - break - elif return_when == ALL_COMPLETED and len(retrying_pending) == 0: - break + retrying_done, retrying_pending = self._split_done_and_retried( + done, pending, lookup + ) - return retrying_done, retrying_pending + if self._wait_is_over( + return_when, retrying_done, retrying_pending + ): + return retrying_done, retrying_pending def clean( self, @@ -347,14 +372,16 @@ def clean( force: Optional[bool] = False ): """ - Cleans up temporary files and objects related to this executor, including: + Cleans up temporary files and objects related to this + executor, including: - Function packages - Serialized input/output data - Cloud objects (if specified) :param fs: List of futures to clean. :param cs: List of cloudobjects to clean. - :param clean_cloudobjects: Whether to delete all cloudobjects created with this executor. + :param clean_cloudobjects: Whether to delete all + cloudobjects created with this executor. :param clean_fn: Whether to delete cached functions. :param force: Force cleanup even for unfinished futures. """ diff --git a/lithops/scripts/cleaner.py b/lithops/scripts/cleaner.py index 43c06bd0a..38986a6b0 100644 --- a/lithops/scripts/cleaner.py +++ b/lithops/scripts/cleaner.py @@ -20,156 +20,352 @@ import pickle import logging from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, Iterable, List, Optional, Tuple from lithops.storage import Storage from lithops.storage.utils import clean_bucket from lithops.constants import JOBS_PREFIX, TEMP_PREFIX, CLEANER_DIR, \ - CLEANER_PID_FILE, CLEANER_LOG_FILE - -log_file_stream = open(CLEANER_LOG_FILE, 'a') -sys.stdout = log_file_stream -sys.stderr = log_file_stream + CLEANER_PID_FILE, CLEANER_LOG_FILE, CLEANER_TMP_SUFFIX +try: + import fcntl +except ImportError: # Windows + fcntl = None +try: + import msvcrt +except ImportError: # everything but Windows + msvcrt = None logger = logging.getLogger('lithops') -logging.basicConfig(stream=log_file_stream, level=logging.INFO, - format=('%(asctime)s [%(levelname)s] %(module)s' - ' [%(threadName)s] - %(funcName)s: %(message)s')) -logger.setLevel('DEBUG') - - -def clean_executor_jobs(executor_id, executor_data): +_SKIP_FILES = {CLEANER_LOG_FILE, CLEANER_PID_FILE} + +# After the last request, wait once more so a parallel job can drop its file +# before this process exits and releases the lock +_IDLE_CONFIRM_SECONDS = 2 + +# A cleaner that finds the lock taken keeps trying for this long, but only +# while there is something pending: the cleaner holding it may be exiting +_LOCK_RETRY_SECONDS = 10 +_LOCK_RETRY_INTERVAL = 0.5 + +# One pickled request dropped in CLEANER_DIR: where the file lives on disk +# and the payload describing what it asks to be deleted +CleanerEntry = Dict[str, Any] + + +def _configure_cleaner_logging() -> None: + """Redirect process output into the cleaner log (subprocess only).""" + os.makedirs(CLEANER_DIR, exist_ok=True) + log_file_stream = open(CLEANER_LOG_FILE, 'a') + sys.stdout = log_file_stream + sys.stderr = log_file_stream + logging.basicConfig( + stream=log_file_stream, + level=logging.INFO, + format=( + '%(asctime)s [%(levelname)s] %(module)s [%(threadName)s] - ' + '%(funcName)s: %(message)s' + ), + ) + logger.setLevel(logging.DEBUG) + + +def _remove_if_exists(path: str) -> None: + try: + os.remove(path) + except FileNotFoundError: + pass + + +def _lock_pid_file() -> Optional[int]: + """ + Takes the machine wide cleaner lock, returning the open descriptor of the + pid file, or None when another cleaner already holds it. + + The lock belongs to the process, not to the contents of the file: the + operating system drops it as soon as this process ends, so a cleaner that + is killed cannot block the next one, and there is no stale pid to detect. + The pid is written for diagnostics only, and the file is never removed: + unlinking it would leave this process holding a lock on an inode nobody + can see, and let a second cleaner lock a fresh file at the same path. + """ + os.makedirs(CLEANER_DIR, exist_ok=True) + flags = os.O_CREAT | os.O_RDWR | getattr(os, 'O_BINARY', 0) + pid_fd = os.open(CLEANER_PID_FILE, flags) + try: + if fcntl is not None: + fcntl.flock(pid_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + elif msvcrt is not None: + msvcrt.locking(pid_fd, msvcrt.LK_NBLCK, 1) + # On a platform with neither there is no lock to take, and every + # cleaner runs. They honour the same requests, which is wasteful + # but not wrong: each one is removed by whoever gets to it first + except OSError: + os.close(pid_fd) + return None + + try: + os.truncate(pid_fd, 0) + os.write(pid_fd, str(os.getpid()).encode()) + except OSError: + # The pid is only there to be read by a human. The lock is taken + # either way, and failing to note it down must not end the cleaner + logger.debug('Could not write the cleaner pid', exc_info=True) + return pid_fd + + +def _acquire_cleaner_lock() -> Optional[int]: + """ + Takes the cleaner lock, retrying for a short while when it is held and + there are requests pending. The cleaner holding it normally honours them, + but it may be on its way out, and then nobody else would pick them up + """ + deadline = time.monotonic() + _LOCK_RETRY_SECONDS + while True: + pid_fd = _lock_pid_file() + if pid_fd is not None: + return pid_fd + if not _pending_request_files() or time.monotonic() >= deadline: + return None + time.sleep(_LOCK_RETRY_INTERVAL) + + +def _clean_job_prefixes( + storage: Storage, root_prefix: str, job_keys: Iterable[str], what: str +) -> None: + """ + Deletes everything each job stored under the given top level prefix + """ + for job_key in job_keys: + prefix = '/'.join([root_prefix, job_key]) + '/' + logger.debug(f"Cleaning {what} from {prefix}") + clean_bucket(storage, storage.bucket, prefix) + + +def clean_executor_jobs( + executor_id: str, executor_data: List[CleanerEntry] +) -> None: + """ + Deletes the data left behind by the jobs of a single executor. Every + entry targets the same executor, so one Storage client serves them all + """ storage = None - logger.debug(f"Cleaning Executor ID: {executor_id}") for file_data in executor_data: file_location = file_data['file_location'] data = file_data['data'] - storage_config = data['storage_config'] - clean_cloudobjects = data['clean_cloudobjects'] - logger.debug(f"File location: {file_location}") - if not storage: - storage = Storage(storage_config=storage_config) - - for job_key in data['jobs_to_clean']: - prefix = '/'.join([JOBS_PREFIX, job_key]) + '/' - logger.debug(f"Cleaning data from {prefix}") - clean_bucket(storage, storage.bucket, prefix) + if storage is None: + storage = Storage(storage_config=data['storage_config']) - if clean_cloudobjects: - for job_key in data['jobs_to_clean']: - prefix = '/'.join([TEMP_PREFIX, job_key]) + '/' - logger.debug(f"Cleaning cloudobjects from {prefix}") - clean_bucket(storage, storage.bucket, prefix) + _clean_job_prefixes( + storage, JOBS_PREFIX, data['jobs_to_clean'], 'data' + ) + if data['clean_cloudobjects']: + _clean_job_prefixes( + storage, TEMP_PREFIX, data['jobs_to_clean'], 'cloudobjects' + ) - if os.path.exists(file_location): - os.remove(file_location) + _remove_if_exists(file_location) logger.info('Finished') -def clean_cloudobjects(cloudobjects_data): +def clean_cloudobjects(cloudobjects_data: CleanerEntry) -> None: + """ + Deletes the cloudobjects of a request, skipping the ones that live in a + storage backend other than the one the request was created with + """ file_location = cloudobjects_data['file_location'] data = cloudobjects_data['data'] logger.info('Going to clean cloudobjects') - cos_to_clean = data['cos_to_clean'] - storage_config = data['storage_config'] - storage = Storage(storage_config=storage_config) + storage = Storage(storage_config=data['storage_config']) - for co in cos_to_clean: + for co in data['cos_to_clean']: if co.backend == storage.backend: - logging.info('Cleaning {}://{}/{}'.format(co.backend, - co.bucket, - co.key)) + logger.info(f'Cleaning {co.backend}://{co.bucket}/{co.key}') storage.delete_object(co.bucket, co.key) - if os.path.exists(file_location): - os.remove(file_location) + _remove_if_exists(file_location) logger.info('Finished') -def clean_functions(functions_data): +def clean_functions(functions_data: CleanerEntry) -> None: + """ + Deletes the serialized functions an executor uploaded + """ file_location = functions_data['file_location'] data = functions_data['data'] - executor_id = data['fn_to_clean'] - storage_config = data['storage_config'] - storage = Storage(storage_config=storage_config) - prefix = '/'.join([JOBS_PREFIX, executor_id]) + '/' + storage = Storage(storage_config=data['storage_config']) + prefix = '/'.join([JOBS_PREFIX, data['fn_to_clean']]) + '/' logger.info(f'Cleaning functions from {prefix}') key_list = storage.list_keys(storage.bucket, prefix) storage.delete_objects(storage.bucket, key_list) - if os.path.exists(file_location): - os.remove(file_location) + _remove_if_exists(file_location) logger.info('Finished') -def clean(): - - while True: - executor_jobs = {} - cloudobjects = [] - functions = [] - - files_to_clean = os.listdir(CLEANER_DIR) - - if len(files_to_clean) <= 2: - break +def _load_cleaner_file(file_location: str) -> Dict[str, Any]: + with open(file_location, 'rb') as pk: + return pickle.load(pk) - for file_name in files_to_clean: - file_location = os.path.join(CLEANER_DIR, file_name) - if file_location in [CLEANER_LOG_FILE, CLEANER_PID_FILE]: - continue - with open(file_location, 'rb') as pk: - data = pickle.load(pk) +def _executor_id_from_jobs(jobs_to_clean: Iterable[str]) -> Optional[str]: + """ + Derives the executor id of a job key, which is the key minus its job + number. Returns None when there is no job to derive it from + """ + first_key = next(iter(jobs_to_clean), None) + if not first_key: + return None + return first_key.rsplit('-', 1)[0] - if 'jobs_to_clean' in data: - # group data by executor_id - executor_id, job_id = next(iter(data['jobs_to_clean'])).rsplit('-', 1) - if executor_id not in executor_jobs: - executor_jobs[executor_id] = [] - executor_jobs[executor_id].append({'file_location': file_location, 'data': data}) - elif 'cos_to_clean' in data: - cloudobjects.append({'file_location': file_location, 'data': data}) +def _classify_cleaner_files( + files_to_clean: List[str] +) -> Tuple[Dict[str, List[CleanerEntry]], List[CleanerEntry], List[CleanerEntry]]: + """ + Reads every pending request and sorts it by the kind of data it asks to + delete, grouping the job requests by executor so that they share a client + """ + executor_jobs: Dict[str, List[CleanerEntry]] = {} + cloudobjects: List[CleanerEntry] = [] + functions: List[CleanerEntry] = [] - elif 'fn_to_clean' in data: - functions.append({'file_location': file_location, 'data': data}) + for file_name in files_to_clean: + file_location = os.path.join(CLEANER_DIR, file_name) + if file_location in _SKIP_FILES: + continue - if executor_jobs: - with ThreadPoolExecutor(max_workers=32) as ex: - for executor_id in executor_jobs: - ex.submit(clean_executor_jobs, executor_id, executor_jobs[executor_id]) + try: + data = _load_cleaner_file(file_location) + except FileNotFoundError: + # Honoured between the listing and the read + continue + except Exception: + # A request nobody can read is dropped rather than retried: it + # would otherwise be picked up on every pass, and the loop below + # would never see an empty directory again + logger.warning( + f'Discarding unreadable request {file_location}', + exc_info=True + ) + _remove_if_exists(file_location) + continue + + entry = {'file_location': file_location, 'data': data} + + if not isinstance(data, dict): + logger.warning(f'Discarding {file_location}: not a request') + _remove_if_exists(file_location) + elif 'jobs_to_clean' in data: + executor_id = _executor_id_from_jobs(data['jobs_to_clean']) + if executor_id is None: + logger.warning( + f'Skipping {file_location}: jobs_to_clean is empty' + ) + _remove_if_exists(file_location) + continue + executor_jobs.setdefault(executor_id, []).append(entry) + elif 'cos_to_clean' in data: + cloudobjects.append(entry) + elif 'fn_to_clean' in data: + functions.append(entry) + else: + logger.warning( + f'Discarding {file_location}: unknown request {sorted(data)}' + ) + _remove_if_exists(file_location) + + return executor_jobs, cloudobjects, functions + + +def _run_clean_tasks( + executor_jobs: Dict[str, List[CleanerEntry]], + cloudobjects: List[CleanerEntry], + functions: List[CleanerEntry], +) -> None: + """ + Runs every classified request in parallel and waits for all of them, + re-raising the first failure so that it is not silently swallowed + """ + tasks = [] + with ThreadPoolExecutor(max_workers=32) as ex: + for executor_id, jobs in executor_jobs.items(): + tasks.append(ex.submit(clean_executor_jobs, executor_id, jobs)) + for item in cloudobjects: + tasks.append(ex.submit(clean_cloudobjects, item)) + for item in functions: + tasks.append(ex.submit(clean_functions, item)) + for task in tasks: + task.result() + + +def _pending_request_files() -> List[str]: + """ + Lists the request files waiting in CLEANER_DIR, leaving out the log and + the pid file, which are not requests and may or may not be there, and + the staging files of a request another process is still writing + """ + try: + file_names = os.listdir(CLEANER_DIR) + except FileNotFoundError: + return [] + + return [ + file_name for file_name in file_names + if not file_name.endswith(CLEANER_TMP_SUFFIX) + and os.path.join(CLEANER_DIR, file_name) not in _SKIP_FILES + ] + + +def clean() -> None: + """ + Processes the pending requests until none is left. After a quiet poll, + wait once more so a parallel Lithops command can still drop a request + before this process exits. + """ + while True: + files_to_clean = _pending_request_files() + if not files_to_clean: + time.sleep(_IDLE_CONFIRM_SECONDS) + files_to_clean = _pending_request_files() + if not files_to_clean: + break + + executor_jobs, cloudobjects, functions = _classify_cleaner_files( + files_to_clean + ) + _run_clean_tasks(executor_jobs, cloudobjects, functions) + time.sleep(5) - if cloudobjects: - with ThreadPoolExecutor(max_workers=32) as ex: - for cloudobjects_data in cloudobjects: - ex.submit(clean_cloudobjects, cloudobjects_data) - if functions: - with ThreadPoolExecutor(max_workers=32) as ex: - for function_data in functions: - ex.submit(clean_functions, function_data) +def main() -> None: + """ + Entry point of the cleaner process. One cleaner serves every Lithops + command on this machine, so a run that cannot take the lock exits and + lets the one holding it honour the requests of them all + """ + pid_fd = _acquire_cleaner_lock() + if pid_fd is None: + return - time.sleep(5) + _configure_cleaner_logging() + logger.info("Starting Job and Cloudobject Cleaner") + try: + clean() + finally: + # Closing drops the lock, which is what lets the next cleaner start. + # The pid file itself stays, see _lock_pid_file() + os.close(pid_fd) + logger.info("Job and Cloudobject Cleaner finished") if __name__ == '__main__': - if not os.path.isfile(CLEANER_PID_FILE): - logger.info("Starting Job and Cloudobject Cleaner") - with open(CLEANER_PID_FILE, 'w') as cf: - cf.write(str(os.getpid())) - try: - clean() - except Exception as e: - raise e - finally: - os.remove(CLEANER_PID_FILE) - logger.info("Job and Cloudobject Cleaner finished") + main() diff --git a/lithops/scripts/cli.py b/lithops/scripts/cli.py index 7fe835f6c..f329c4987 100644 --- a/lithops/scripts/cli.py +++ b/lithops/scripts/cli.py @@ -15,18 +15,18 @@ # limitations under the License. # - import os import time import click +import getpass import logging import shutil -import shlex import subprocess as sp from itertools import cycle from tabulate import tabulate from datetime import datetime from concurrent.futures import ThreadPoolExecutor +from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union import lithops from lithops import Storage @@ -47,6 +47,8 @@ ) from lithops.constants import ( CACHE_DIR, + CLEANER_DIR, + JOBS_DIR, LITHOPS_TEMP_DIR, RUNTIMES_PREFIX, JOBS_PREFIX, @@ -67,25 +69,307 @@ logger = logging.getLogger(__name__) -def set_config_ow(backend, storage=None, runtime_name=None, region=None): +def set_config_ow( + backend: Optional[str] = None, + storage: Optional[str] = None, + runtime_name: Optional[str] = None, + region: Optional[str] = None, +) -> Dict[str, Any]: + """ + Builds the config overwrite that the global CLI options impose on top of + whatever the user config file provides + """ config_ow = {'lithops': {}, 'backend': {}} - if storage: config_ow['lithops']['storage'] = storage - if backend: config_ow['lithops']['backend'] = backend config_ow['lithops']['mode'] = get_mode(backend) - if runtime_name: config_ow['backend']['runtime'] = runtime_name - if region: config_ow['backend']['region'] = region - return config_ow +def _load_user_config(config_path: Optional[str]) -> Optional[Dict[str, Any]]: + return load_yaml_config(config_path) if config_path else None + + +def _setup_cli_logger(debug: bool) -> None: + setup_lithops_logger(logging.DEBUG if debug else logging.INFO) + + +def _resolved_config( + config_path: Optional[str], + *, + backend: Optional[str] = None, + storage: Optional[str] = None, + runtime_name: Optional[str] = None, + region: Optional[str] = None, + load_storage_config: bool = True, +) -> Dict[str, Any]: + """ + Loads the user config file, if any, and merges the CLI options into it + """ + config = _load_user_config(config_path) + config_ow = set_config_ow( + backend=backend, + storage=storage, + runtime_name=runtime_name, + region=region, + ) + return default_config( + config_data=config, + config_overwrite=config_ow, + load_storage_config=load_storage_config, + ) + + +def _require_mode(config: Dict[str, Any], mode: str, command: str) -> None: + """ + Rejects a command that the configured compute mode cannot serve + """ + if config['lithops']['mode'] == mode: + return + if mode == STANDALONE: + raise Exception( + f'{command} is only available for standalone backends. ' + f'Please use "{command} -b {set(STANDALONE_BACKENDS)}"' + ) + raise Exception( + f'"{command}" command is only available for serverless backends' + ) + + +def _standalone_handler(config: Dict[str, Any]) -> StandaloneHandler: + return StandaloneHandler(extract_standalone_config(config)) + + +def _serverless_handler( + config: Dict[str, Any], internal_storage: Optional[InternalStorage] = None +) -> ServerlessHandler: + return ServerlessHandler( + extract_serverless_config(config), internal_storage + ) + + +def _compute_handler( + config: Dict[str, Any], internal_storage: Optional[InternalStorage] = None +): + """ + Builds the compute handler that matches the configured compute mode + """ + mode = config['lithops']['mode'] + if mode == LOCALHOST: + return LocalhostHandler(extract_localhost_config(config)) + if mode == SERVERLESS: + return _serverless_handler(config, internal_storage) + if mode == STANDALONE: + return _standalone_handler(config) + raise Exception(f'Unknown compute mode: {mode}') + + +def _prepare_serverless( + name: Optional[str], + config_path: Optional[str], + backend: Optional[str], + storage: Optional[str], + debug: bool, + command: str, + *, + always_debug: bool = False, + load_storage: bool = True, +) -> Tuple[ServerlessHandler, Optional[InternalStorage]]: + """ + Common setup of the serverless commands: logging, runtime name checks, + config resolution and the handler the command then drives + """ + _setup_cli_logger(debug or always_debug) + if name: + verify_runtime_name(name) + config = _resolved_config( + config_path, + backend=backend, + storage=storage, + runtime_name=name, + load_storage_config=load_storage, + ) + _require_mode(config, SERVERLESS, command) + internal_storage = ( + InternalStorage(extract_storage_config(config)) if load_storage else None + ) + return _serverless_handler(config, internal_storage), internal_storage + + +def _prepare_standalone( + config_path: Optional[str], + backend: Optional[str], + region: Optional[str], + debug: bool, + command: str, + *, + always_debug: bool = False, +) -> StandaloneHandler: + """ + Common setup of the standalone commands. Standalone never needs the + object storage, so it is not loaded + """ + _setup_cli_logger(debug or always_debug) + config = _resolved_config( + config_path, + backend=backend, + region=region, + load_storage_config=False, + ) + _require_mode(config, STANDALONE, command) + return _standalone_handler(config) + + +def _standalone_service_ready(handler: StandaloneHandler) -> bool: + """ + Tells whether the master VM is up and serving, logging why it is not + """ + if not handler.is_initialized(): + logger.info("The backend is not initialized") + return False + handler.init() + if not handler.backend.master.is_ready(): + logger.info(f"{handler.backend.master} is stopped") + return False + if not handler._is_master_service_ready(): + logger.info( + f"Lithops service is not running in {handler.backend.master}" + ) + return False + return True + + +def _utc_to_local(utc_timestamp: str, local_tz) -> str: + import pytz + utc_time = datetime.strptime(utc_timestamp, '%Y-%m-%d %H:%M:%S %Z') + utc_time = utc_time.replace(tzinfo=pytz.utc) + local_time = utc_time.astimezone(local_tz) + return local_time.strftime('%Y-%m-%d %H:%M:%S %Z') + + +def _localize_and_sort_rows(rows: List[List], key_index: int) -> List[List]: + """ + Sorts the rows by their timestamp column, rewritten to the local time + zone. pytz and tzlocal are optional, so without them times stay in UTC + """ + try: + from tzlocal import get_localzone + local_tz = get_localzone() + for row in rows: + row[key_index] = _utc_to_local(row[key_index], local_tz) + except ModuleNotFoundError: + pass + return sorted(rows, key=lambda row: row[key_index]) + + +def _print_table( + rows: List, headers: Union[str, List[str]], total_label: str +) -> None: + print() + print(tabulate(rows, headers=headers)) + print(f'\nTotal {total_label}: {len(rows)}') + + +def _format_storage_objects(objects: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Keeps only the object attributes the storage list command shows, in a + display-friendly form + """ + formatted = [] + for obj in objects: + row = {} + if 'Key' in obj: + row['Key'] = obj['Key'] + if 'LastModified' in obj: + row['LastModified'] = obj['LastModified'].strftime( + "%b %d %Y %H:%M:%S" + ) + if 'Size' in obj: + row['Size'] = sizeof_fmt(obj['Size']) + formatted.append(row) + return formatted + + +def _run_with_spinner(message: str, func: Callable) -> None: + """ + Runs a blocking call on a worker thread while animating a spinner, then + re-raises whatever the call raised + """ + with ThreadPoolExecutor() as ex: + future = ex.submit(func) + spinner = cycle(r"-\|/") + while not future.done(): + print(f"{message} {next(spinner)}", end="\r") + time.sleep(0.1) + future.result() + + +def _make_storage( + config_path: Optional[str], backend: Optional[str], debug: bool +) -> Storage: + _setup_cli_logger(debug) + return Storage(config=_load_user_config(config_path), backend=backend) + + +def _follow_log(fileobj) -> Iterator[str]: + """ + Yields complete lines as they are appended to the log, tail -f style, + and stops once the log file is gone + """ + line = '' + while True: + if not os.path.isfile(FN_LOG_FILE): + break + tmp = fileobj.readline() + if tmp: + line += tmp + if line.endswith("\n"): + yield line + line = '' + else: + time.sleep(1) + + +def _clean_local_temp_data() -> None: + """ + Deletes the local temporary data of this machine: logs, cached modules, + custom runtimes and localhost job data. + + Only the contents go, never the directory skeleton that lithops/config.py + creates on import: a Lithops process that is already running would never + see those directories come back, and would die writing its next log line. + CLEANER_DIR is skipped altogether, since it holds the pending cleaner + requests of every process on this machine plus the lock of the running + cleaner, and dropping a request would leak the data it asks to delete. + Anything else a parallel job is using does go, which is intended: this + command is explicitly destructive. + """ + try: + entries = list(os.scandir(LITHOPS_TEMP_DIR)) + except FileNotFoundError: + entries = [] + + for entry in entries: + if entry.path == CLEANER_DIR: + continue + if entry.is_dir(follow_symlinks=False): + shutil.rmtree(entry.path, ignore_errors=True) + else: + try: + os.remove(entry.path) + except OSError: + pass + + for temp_dir in (LITHOPS_TEMP_DIR, JOBS_DIR, LOGS_DIR, CLEANER_DIR): + os.makedirs(temp_dir, exist_ok=True) + + @click.group('lithops_cli') @click.version_option() def lithops_cli(): @@ -93,65 +377,65 @@ def lithops_cli(): @lithops_cli.command('clean') -@click.option('--config', '-c', default=None, help='path to yaml config file', type=click.Path(exists=True)) +@click.option( + '--config', '-c', default=None, + help='path to yaml config file', type=click.Path(exists=True) +) @click.option('--backend', '-b', default=None, help='compute backend') @click.option('--storage', '-s', default=None, help='storage backend') @click.option('--debug', '-d', is_flag=True, help='debug mode') @click.option('--region', '-r', default=None, help='compute backend region') -@click.option('--all', '-a', is_flag=True, help='delete all, including master VM in case of standalone') -def clean(config, backend, storage, debug, region, all): - config = load_yaml_config(config) if config else None - - log_level = logging.INFO if not debug else logging.DEBUG - setup_lithops_logger(log_level) +@click.option( + '--all', '-a', 'delete_all', is_flag=True, + help='delete all, including master VM in case of standalone' +) +def clean(config, backend, storage, debug, region, delete_all): + _setup_cli_logger(debug) logger.info('Cleaning all Lithops information') - config_ow = set_config_ow(backend=backend, storage=storage, region=region) - config = default_config(config_data=config, config_overwrite=config_ow) - storage_config = extract_storage_config(config) - internal_storage = InternalStorage(storage_config) - - mode = config['lithops']['mode'] - backend = config['lithops']['backend'] - - if mode == LOCALHOST: - compute_config = extract_localhost_config(config) - compute_handler = LocalhostHandler(compute_config) - elif mode == SERVERLESS: - compute_config = extract_serverless_config(config) - compute_handler = ServerlessHandler(compute_config, internal_storage) - elif mode == STANDALONE: - compute_config = extract_standalone_config(config) - compute_handler = StandaloneHandler(compute_config) - - compute_handler.clean(all=all) - - # Clean object storage temp dirs - storage = internal_storage.storage - runtimes_path = RUNTIMES_PREFIX + '/' + backend - jobs_path = JOBS_PREFIX - clean_bucket(storage, storage.bucket, runtimes_path, sleep=1) - clean_bucket(storage, storage.bucket, jobs_path, sleep=1) - - # Clean localhost executor temp dirs - shutil.rmtree(LITHOPS_TEMP_DIR, ignore_errors=True) - # Clean local lithops runtime cache - shutil.rmtree(os.path.join(CACHE_DIR, RUNTIMES_PREFIX, backend), ignore_errors=True) - + cfg = _resolved_config( + config, backend=backend, storage=storage, region=region + ) + compute_backend = cfg['lithops']['backend'] + internal_storage = InternalStorage(extract_storage_config(cfg)) + compute_handler = _compute_handler(cfg, internal_storage) + compute_handler.clean(all=delete_all) + + obj_storage = internal_storage.storage + runtimes_path = f'{RUNTIMES_PREFIX}/{compute_backend}' + clean_bucket(obj_storage, obj_storage.bucket, runtimes_path, sleep=1) + clean_bucket(obj_storage, obj_storage.bucket, JOBS_PREFIX, sleep=1) + + _clean_local_temp_data() + shutil.rmtree( + os.path.join(CACHE_DIR, RUNTIMES_PREFIX, compute_backend), + ignore_errors=True, + ) logger.info('All Lithops temporary data cleaned') @lithops_cli.command('test') -@click.option('--config', '-c', default=None, help='Path to yaml config file', type=click.Path(exists=True)) +@click.option( + '--config', '-c', default=None, + help='Path to yaml config file', type=click.Path(exists=True) +) @click.option('--backend', '-b', default=None, help='Compute backend') @click.option('--storage', '-s', default=None, help='Storage backend') @click.option('--debug', '-d', is_flag=True, help='Debug mode') @click.option('--region', '-r', default=None, help='compute backend region') -@click.option('--test', '-t', default=None, help='Run a specific test. To avoid running similarly named tests ' - 'you may prefix the tester with its test class, ' - 'e.g. TestAsync::test_call_async' - 'Type "-t help" for the complete tests list') -@click.option('--exitfirst', '-x', is_flag=True, help='Stops test run upon first occurrence of a failed test') +@click.option( + '--test', '-t', default=None, + help=( + 'Run a specific test. To avoid running similarly named tests ' + 'you may prefix the tester with its test class, ' + 'e.g. TestAsync::test_call_async ' + 'Type "-t help" for the complete tests list' + ) +) +@click.option( + '--exitfirst', '-x', is_flag=True, + help='Stops test run upon first occurrence of a failed test' +) def test(test, config, backend, storage, debug, region, exitfirst): import pytest @@ -160,59 +444,61 @@ def test(test, config, backend, storage, debug, region, exitfirst): if test == 'help': pytest.main([tests_path, "--collect-only"]) - else: - cmd_string = [tests_path, "-v"] - if exitfirst: - cmd_string.extend(["-x"]) - if debug: - cmd_string.extend(["-o", "log_cli=true", "--log-cli-level=DEBUG"]) - if config: - cmd_string.extend(["--config", config]) - if backend: - cmd_string.extend(["--backend", backend]) - if storage: - cmd_string.extend(["--storage", storage]) - if region: - cmd_string.extend(["--region", region]) - if test: - cmd_string.extend(["-k", test]) - - print("Executing lithops tests: pytest " + ' '.join(cmd_string[1:])) - - pytest.main(cmd_string) + return + + cmd = [tests_path, "-v"] + if exitfirst: + cmd.append("-x") + if debug: + cmd.extend(["-o", "log_cli=true", "--log-cli-level=DEBUG"]) + for option, value in ( + ("--config", config), + ("--backend", backend), + ("--storage", storage), + ("--region", region), + ("-k", test), + ): + if value: + cmd.extend([option, value]) + + print(f"Executing lithops tests: pytest {' '.join(cmd[1:])}") + pytest.main(cmd) @lithops_cli.command('hello') -@click.option('--config', '-c', default=None, help='path to yaml config file', type=click.Path(exists=True)) +@click.option( + '--config', '-c', default=None, + help='path to yaml config file', type=click.Path(exists=True) +) @click.option('--backend', '-b', default=None, help='compute backend') @click.option('--storage', '-s', default=None, help='storage backend') @click.option('--debug', '-d', is_flag=True, help='debug mode') @click.option('--region', '-r', default=None, help='compute backend region') -@click.option('--map', 'map_count', '-m', default=None, type=click.IntRange(min=1), - help='number of map invocations to run instead of a single call_async') +@click.option( + '--map', 'map_count', '-m', default=None, type=click.IntRange(min=1), + help='number of map invocations to run instead of a single call_async' +) def hello(config, backend, storage, debug, region, map_count): - config = load_yaml_config(config) if config else None - - log_level = logging.INFO if not debug else logging.DEBUG - setup_lithops_logger(log_level) + _setup_cli_logger(debug) + config_data = _load_user_config(config) try: - import getpass username = getpass.getuser() - except Exception: + except (OSError, KeyError): + # No login name and no matching passwd entry, which happens in some + # containers username = 'World' - def hello(name): + def hello_fn(name): return f'Hello {name}!' fexec = lithops.FunctionExecutor( - config=config, backend=backend, - storage=storage, region=region + config=config_data, backend=backend, storage=storage, region=region ) expected = f'Hello {username}!' if map_count: - fexec.map(hello, [username] * map_count) + fexec.map(hello_fn, [username] * map_count) results = fexec.get_result() print() if all(result == expected for result in results): @@ -221,7 +507,7 @@ def hello(name): else: print(results, 'Something went wrong :(') else: - fexec.call_async(hello, username) + fexec.call_async(hello_fn, username) result = fexec.get_result() print() if result == expected: @@ -232,56 +518,55 @@ def hello(name): @lithops_cli.command('attach') -@click.option('--config', '-c', default=None, help='path to yaml config file', type=click.Path(exists=True)) +@click.option( + '--config', '-c', default=None, + help='path to yaml config file', type=click.Path(exists=True) +) @click.option('--backend', '-b', default=None, help='compute backend') @click.option("--start", is_flag=True, default=False, help="Start the master VM if needed.") @click.option('--debug', '-d', is_flag=True, help='debug mode') @click.option('--region', '-r', default=None, help='compute backend region') def attach(config, backend, start, debug, region): """Create or attach to a SSH session on Lithops master VM""" - config = load_yaml_config(config) if config else None - - log_level = logging.INFO if not debug else logging.DEBUG - setup_lithops_logger(log_level) - - config_ow = set_config_ow(backend=backend, region=region) - config = default_config(config_data=config, config_overwrite=config_ow, load_storage_config=False) - - if config['lithops']['mode'] != STANDALONE: - raise Exception('lithops attach method is only available for standalone backends. ' - f'Please use "lithops attach -b {set(STANDALONE_BACKENDS)}"') - - compute_config = extract_standalone_config(config) - compute_handler = StandaloneHandler(compute_config) + handler = _prepare_standalone( + config, backend, region, debug, 'lithops attach' + ) - if not compute_handler.is_initialized(): + if not handler.is_initialized(): logger.info("The backend is not initialized") return - compute_handler.init() - if not start and not compute_handler.backend.master.is_ready(): - logger.info(f"{compute_handler.backend.master} is stopped") + handler.init() + if not start and not handler.backend.master.is_ready(): + logger.info(f"{handler.backend.master} is stopped") return if start: - compute_handler.backend.master.start() + handler.backend.master.start() - master_ip = compute_handler.backend.master.get_public_ip() - user = compute_handler.backend.master.ssh_credentials['username'] - key_file = compute_handler.backend.master.ssh_credentials['key_filename'] or '~/.ssh/id_rsa' + master_ip = handler.backend.master.get_public_ip() + user = handler.backend.master.ssh_credentials['username'] + key_file = ( + handler.backend.master.ssh_credentials['key_filename'] + or '~/.ssh/id_rsa' + ) key_file = os.path.abspath(os.path.expanduser(key_file)) if not os.path.exists(key_file): - raise Exception(f'Private key file {key_file} does not exists') + raise Exception(f'Private key file {key_file} does not exist') print(f'Got master VM public IP address: {master_ip}') print(f'Loading ssh private key from: {key_file}') print('Creating SSH Connection to lithops master VM') - cmd = ('ssh -o "UserKnownHostsFile=/dev/null" -o "StrictHostKeyChecking=no" ' - f'-i {key_file} {user}@{master_ip}') - - compute_handler.backend.master.wait_ready() + cmd = [ + 'ssh', + '-o', 'UserKnownHostsFile=/dev/null', + '-o', 'StrictHostKeyChecking=no', + '-i', key_file, + f'{user}@{master_ip}', + ] - sp.run(shlex.split(cmd)) + handler.backend.master.wait_ready() + sp.run(cmd) # /---------------------------------------------------------------------------/ @@ -291,8 +576,7 @@ def attach(config, backend, start, debug, region): # /---------------------------------------------------------------------------/ @click.group('storage') -@click.pass_context -def storage(ctx): +def storage(): pass @@ -302,29 +586,28 @@ def storage(ctx): @click.option('--key', '-k', default=None, help='object key') @click.option('--backend', '-b', default=None, help='storage backend') @click.option('--debug', '-d', is_flag=True, help='debug mode') -@click.option('--config', '-c', default=None, help='path to yaml config file', type=click.Path(exists=True)) +@click.option( + '--config', '-c', default=None, + help='path to yaml config file', type=click.Path(exists=True) +) def upload_file(filename, bucket, key, backend, debug, config): - config = load_yaml_config(config) if config else None - - log_level = logging.INFO if not debug else logging.DEBUG - setup_lithops_logger(log_level) - storage = Storage(config=config, backend=backend) - - def upload_file(): - logger.info(f'Uploading file {filename} to {storage.backend}://{bucket}/{key or filename}') - if storage.upload_file(filename, bucket, key): + client = _make_storage(config, backend, debug) + dest_key = key or filename + + def _upload(): + logger.info( + f'Uploading file {filename} to ' + f'{client.backend}://{bucket}/{dest_key}' + ) + if client.upload_file(filename, bucket, key): file_size = os.path.getsize(filename) - logger.info(f'Upload File {filename} - Size: {sizeof_fmt(file_size)} - Ok') + logger.info( + f'Upload File {filename} - Size: {sizeof_fmt(file_size)} - Ok' + ) else: logger.error(f'Upload File {filename} - Error') - with ThreadPoolExecutor() as ex: - future = ex.submit(upload_file) - cy = cycle(r"-\|/") - while not future.done(): - print("Uploading file " + next(cy), end="\r") - time.sleep(0.1) - future.result() + _run_with_spinner("Uploading file", _upload) @storage.command('get') @@ -333,29 +616,27 @@ def upload_file(): @click.option('--out', '-o', default=None, help='output filename') @click.option('--backend', '-b', default=None, help='storage backend') @click.option('--debug', '-d', is_flag=True, help='debug mode') -@click.option('--config', '-c', default=None, help='path to yaml config file', type=click.Path(exists=True)) +@click.option( + '--config', '-c', default=None, + help='path to yaml config file', type=click.Path(exists=True) +) def download_file(bucket, key, out, backend, debug, config): - config = load_yaml_config(config) if config else None - - log_level = logging.INFO if not debug else logging.DEBUG - setup_lithops_logger(log_level) - storage = Storage(config=config, backend=backend) - - def download_file(): - logger.info(f'Downloading file {storage.backend}://{bucket}/{key} to {out or key}') - if storage.download_file(bucket, key, out): - file_size = os.path.getsize(out or key) - logger.info(f'Download File {key} - Size: {sizeof_fmt(file_size)} - Ok') + client = _make_storage(config, backend, debug) + dest = out or key + + def _download(): + logger.info( + f'Downloading file {client.backend}://{bucket}/{key} to {dest}' + ) + if client.download_file(bucket, key, out): + file_size = os.path.getsize(dest) + logger.info( + f'Download File {key} - Size: {sizeof_fmt(file_size)} - Ok' + ) else: logger.error(f'Download File {key} - Error') - with ThreadPoolExecutor() as ex: - future = ex.submit(download_file) - cy = cycle(r"-\|/") - while not future.done(): - print("Downloading file " + next(cy), end="\r") - time.sleep(0.1) - future.result() + _run_with_spinner("Downloading file", _download) @storage.command('delete') @@ -364,22 +645,30 @@ def download_file(): @click.option('--prefix', '-p', default=None, help='key prefix') @click.option('--backend', '-b', default=None, help='storage backend') @click.option('--debug', '-d', is_flag=True, help='debug mode') -@click.option('--config', '-c', default=None, help='path to yaml config file', type=click.Path(exists=True)) +@click.option( + '--config', '-c', default=None, + help='path to yaml config file', type=click.Path(exists=True) +) def delete_object(bucket, key, prefix, backend, debug, config): - config = load_yaml_config(config) if config else None - log_level = logging.INFO if not debug else logging.DEBUG - setup_lithops_logger(log_level) - storage = Storage(config=config, backend=backend) + client = _make_storage(config, backend, debug) if key: - logger.info('Deleting object "{}" from bucket "{}"'.format(key, bucket)) - storage.delete_object(bucket, key) - logger.info('Object deleted successfully') - elif prefix: - objs = storage.list_keys(bucket, prefix) - logger.info('Deleting {} objects with prefix "{}" from bucket "{}"'.format(len(objs), prefix, bucket)) - storage.delete_objects(bucket, objs) + logger.info(f'Deleting object "{key}" from bucket "{bucket}"') + client.delete_object(bucket, key) logger.info('Object deleted successfully') + return + + if prefix: + objs = client.list_keys(bucket, prefix) + logger.info( + f'Deleting {len(objs)} objects with prefix "{prefix}" ' + f'from bucket "{bucket}"' + ) + client.delete_objects(bucket, objs) + logger.info('Objects deleted successfully') + return + + raise click.UsageError('Provide KEY or --prefix') @storage.command('list') @@ -387,30 +676,22 @@ def delete_object(bucket, key, prefix, backend, debug, config): @click.option('--prefix', '-p', default=None, help='key prefix') @click.option('--backend', '-b', default=None, help='storage backend') @click.option('--debug', '-d', is_flag=True, help='debug mode') -@click.option('--config', '-c', default=None, help='path to yaml config file', type=click.Path(exists=True)) +@click.option( + '--config', '-c', default=None, + help='path to yaml config file', type=click.Path(exists=True) +) def list_bucket(prefix, bucket, backend, debug, config): - config = load_yaml_config(config) if config else None - log_level = logging.INFO if not debug else logging.DEBUG - setup_lithops_logger(log_level) - storage = Storage(config=config, backend=backend) - logger.info('Listing objects in bucket {}'.format(bucket)) - objects = storage.list_objects(bucket, prefix=prefix) - - objs = [ - { - key: obj[key].strftime("%b %d %Y %H:%M:%S") if key == 'LastModified' else sizeof_fmt(obj[key]) if key == 'Size' else obj[key] - for key in ('Key', 'LastModified', 'Size') - if key in obj - } - for obj in objects - ] + client = _make_storage(config, backend, debug) + logger.info(f'Listing objects in bucket {bucket}') + objs = _format_storage_objects(client.list_objects(bucket, prefix=prefix)) - if objs[0]: - print() - print(tabulate(objs, headers="keys")) - print(f'\nTotal objects: {len(objs)}') + if objs: + _print_table(objs, headers="keys", total_label='objects') else: - print(f'\nNo information can be listed from bucket \"{bucket}\" using current \"{storage.backend}\" backend') + print( + f'\nNo information can be listed from bucket "{bucket}" ' + f'using current "{client.backend}" backend' + ) # /---------------------------------------------------------------------------/ @@ -420,8 +701,7 @@ def list_bucket(prefix, bucket, backend, debug, config): # /---------------------------------------------------------------------------/ @click.group('logs') -@click.pass_context -def logs(ctx): +def logs(): pass @@ -429,24 +709,11 @@ def logs(ctx): def poll(): logging.basicConfig(level=logging.DEBUG) - def follow(file): - line = '' - while True: - if not os.path.isfile(FN_LOG_FILE): - break - tmp = file.readline() - if tmp: - line += tmp - if line.endswith("\n"): - yield line - line = '' - else: - time.sleep(1) - while True: if os.path.isfile(FN_LOG_FILE): - for line in follow(open(FN_LOG_FILE, 'r')): - print(line, end='') + with open(FN_LOG_FILE, 'r') as log_file: + for line in _follow_log(log_file): + print(line, end='') else: time.sleep(1) @@ -454,10 +721,10 @@ def follow(file): @logs.command('get') @click.argument('job_key') def get_logs(job_key): - log_file = os.path.join(LOGS_DIR, job_key + '.log') + log_file = os.path.join(LOGS_DIR, f'{job_key}.log') if not os.path.isfile(log_file): - print('The execution id: {} does not exists in logs'.format(job_key)) + print(f'The execution id: {job_key} does not exist in logs') return with open(log_file, 'r') as content_file: @@ -471,44 +738,46 @@ def get_logs(job_key): # /---------------------------------------------------------------------------/ @click.group('runtime') -@click.pass_context -def runtime(ctx): +def runtime(): pass -@runtime.command('build', context_settings=dict(ignore_unknown_options=True, allow_extra_args=True)) +@runtime.command( + 'build', + context_settings=dict(ignore_unknown_options=True, allow_extra_args=True) +) @click.argument('name', required=False) -@click.option('--file', '-f', default=None, help='file needed to build the runtime', type=click.Path(exists=True)) -@click.option('--config', '-c', default=None, help='path to yaml config file', type=click.Path(exists=True)) +@click.option( + '--file', '-f', default=None, + help='file needed to build the runtime', type=click.Path(exists=True) +) +@click.option( + '--config', '-c', default=None, + help='path to yaml config file', type=click.Path(exists=True) +) @click.option('--backend', '-b', default=None, help='compute backend') @click.option('--debug', '-d', is_flag=True, help='debug mode') @click.pass_context def build(ctx, name, file, config, backend, debug): """ build a serverless runtime. """ - # log_level = logging.INFO if not debug else logging.DEBUG - setup_lithops_logger(logging.DEBUG) - - verify_runtime_name(name) - - config = load_yaml_config(config) if config else None - config_ow = set_config_ow(backend=backend, runtime_name=name) - config = default_config(config_data=config, config_overwrite=config_ow, load_storage_config=False) - - if config['lithops']['mode'] != SERVERLESS: - raise Exception('"lithops runtime build" command is only available for serverless backends') - - compute_config = extract_serverless_config(config) - compute_handler = ServerlessHandler(compute_config, None) - runtime_info = compute_handler.get_runtime_info() + handler, _ = _prepare_serverless( + name, config, backend, None, debug, + 'lithops runtime build', + always_debug=True, + load_storage=False, + ) + runtime_info = handler.get_runtime_info() runtime_name = runtime_info['runtime_name'] - compute_handler.build_runtime(runtime_name, file, ctx.args) - + handler.build_runtime(runtime_name, file, ctx.args) logger.info('Runtime built') @runtime.command('deploy') @click.argument('name', required=True) -@click.option('--config', '-c', default=None, help='path to yaml config file', type=click.Path(exists=True)) +@click.option( + '--config', '-c', default=None, + help='path to yaml config file', type=click.Path(exists=True) +) @click.option('--backend', '-b', default=None, help='compute backend') @click.option('--storage', '-s', default=None, help='storage backend') @click.option('--memory', default=None, help='memory used by the runtime', type=int) @@ -516,109 +785,85 @@ def build(ctx, name, file, config, backend, debug): @click.option('--debug', '-d', is_flag=True, help='debug mode') def deploy(name, storage, backend, memory, timeout, config, debug): """ deploy a serverless runtime """ - setup_lithops_logger(logging.DEBUG) - - verify_runtime_name(name) - - config = load_yaml_config(config) if config else None - config_ow = set_config_ow(backend=backend, storage=storage, runtime_name=name) - config = default_config(config_data=config, config_overwrite=config_ow) - - if config['lithops']['mode'] != SERVERLESS: - raise Exception('"lithops runtime deploy" command is only available for serverless backends') - - storage_config = extract_storage_config(config) - internal_storage = InternalStorage(storage_config) - compute_config = extract_serverless_config(config) - compute_handler = ServerlessHandler(compute_config, internal_storage) - - runtime_info = compute_handler.get_runtime_info() + handler, internal_storage = _prepare_serverless( + name, config, backend, storage, debug, + 'lithops runtime deploy', + always_debug=True, + ) + runtime_info = handler.get_runtime_info() runtime_name = runtime_info['runtime_name'] runtime_memory = memory or runtime_info['runtime_memory'] runtime_timeout = timeout or runtime_info['runtime_timeout'] - runtime_key = compute_handler.get_runtime_key(runtime_name, runtime_memory, __version__) - runtime_meta = compute_handler.deploy_runtime(runtime_name, runtime_memory, runtime_timeout) + runtime_key = handler.get_runtime_key( + runtime_name, runtime_memory, __version__ + ) + runtime_meta = handler.deploy_runtime( + runtime_name, runtime_memory, runtime_timeout + ) runtime_meta['runtime_timeout'] = runtime_timeout internal_storage.put_runtime_meta(runtime_key, runtime_meta) - logger.info('Runtime deployed') @runtime.command('list') @click.argument('name', default='all', required=False) -@click.option('--config', '-c', default=None, help='path to yaml config file', type=click.Path(exists=True)) +@click.option( + '--config', '-c', default=None, + help='path to yaml config file', type=click.Path(exists=True) +) @click.option('--backend', '-b', default=None, help='compute backend') @click.option('--storage', '-s', default=None, help='storage backend') @click.option('--debug', '-d', is_flag=True, help='debug mode') def list_runtimes(name, config, backend, storage, debug): """ list all deployed serverless runtime. """ - log_level = logging.INFO if not debug else logging.DEBUG - setup_lithops_logger(log_level) - - config = load_yaml_config(config) if config else None - config_ow = set_config_ow(backend=backend) - config = default_config(config_data=config, config_overwrite=config_ow, load_storage_config=False) - - if config['lithops']['mode'] != SERVERLESS: - raise Exception('"lithops runtime list" command is only available for serverless backends') - - compute_config = extract_serverless_config(config) - compute_handler = ServerlessHandler(compute_config, None) - runtimes = compute_handler.list_runtimes(runtime_name=name) - + handler, _ = _prepare_serverless( + None, config, backend, storage, debug, + 'lithops runtime list', + load_storage=False, + ) + runtimes = handler.list_runtimes(runtime_name=name) headers = ['Runtime Name', 'Memory Size', 'Lithops Version', 'Worker Name'] - - print() - print(tabulate(runtimes, headers=headers)) - print(f'\nTotal runtimes: {len(runtimes)}') + _print_table(runtimes, headers, 'runtimes') @runtime.command('update') @click.argument('name', required=True) -@click.option('--config', '-c', default=None, help='path to yaml config file', type=click.Path(exists=True)) +@click.option( + '--config', '-c', default=None, + help='path to yaml config file', type=click.Path(exists=True) +) @click.option('--backend', '-b', default=None, help='compute backend') @click.option('--storage', '-s', default=None, help='storage backend') @click.option('--debug', '-d', is_flag=True, help='debug mode') def update(name, config, backend, storage, debug): """ Update a serverless runtime """ - log_level = logging.INFO if not debug else logging.DEBUG - setup_lithops_logger(log_level) - - verify_runtime_name(name) - - config = load_yaml_config(config) if config else None - config_ow = set_config_ow(backend=backend, storage=storage, runtime_name=name) - config = default_config(config_data=config, config_overwrite=config_ow) - - if config['lithops']['mode'] != SERVERLESS: - raise Exception('"lithops runtime update" command is only available for serverless backends') - - storage_config = extract_storage_config(config) - internal_storage = InternalStorage(storage_config) - compute_config = extract_serverless_config(config) - compute_handler = ServerlessHandler(compute_config, internal_storage) - - runtime_info = compute_handler.get_runtime_info() + handler, internal_storage = _prepare_serverless( + name, config, backend, storage, debug, 'lithops runtime update' + ) + runtime_info = handler.get_runtime_info() runtime_name = runtime_info['runtime_name'] runtime_timeout = runtime_info['runtime_timeout'] logger.info(f'Updating runtime: {runtime_name}') - runtimes = compute_handler.list_runtimes(runtime_name) - - for runtime in runtimes: - if runtime[2] == __version__: - runtime_key = compute_handler.get_runtime_key(runtime[0], runtime[1], runtime[2]) - runtime_meta = compute_handler.deploy_runtime(runtime[0], runtime[1], runtime_timeout) - internal_storage.put_runtime_meta(runtime_key, runtime_meta) + # Rows are (name, memory, version, worker name), see list_runtimes + for rt in handler.list_runtimes(runtime_name): + if rt[2] != __version__: + continue + runtime_key = handler.get_runtime_key(rt[0], rt[1], rt[2]) + runtime_meta = handler.deploy_runtime(rt[0], rt[1], runtime_timeout) + internal_storage.put_runtime_meta(runtime_key, runtime_meta) logger.info('Runtime updated') @runtime.command('delete') @click.argument('name', required=True) -@click.option('--config', '-c', default=None, help='path to yaml config file', type=click.Path(exists=True)) +@click.option( + '--config', '-c', default=None, + help='path to yaml config file', type=click.Path(exists=True) +) @click.option('--memory', '-m', default=None, help='runtime memory') @click.option('--version', '-v', default=None, help='lithops version') @click.option('--backend', '-b', default=None, help='compute backend') @@ -626,45 +871,27 @@ def update(name, config, backend, storage, debug): @click.option('--debug', '-d', is_flag=True, help='debug mode') def delete(name, config, memory, version, backend, storage, debug): """ delete a serverless runtime """ - log_level = logging.INFO if not debug else logging.DEBUG - setup_lithops_logger(log_level) - - verify_runtime_name(name) - - config = load_yaml_config(config) if config else None - config_ow = set_config_ow(backend=backend, storage=storage, runtime_name=name) - config = default_config(config_data=config, config_overwrite=config_ow) - - if config['lithops']['mode'] != SERVERLESS: - raise Exception('"lithops runtime delete" command is only available for serverless backends') - - storage_config = extract_storage_config(config) - internal_storage = InternalStorage(storage_config) - compute_config = extract_serverless_config(config) - compute_handler = ServerlessHandler(compute_config, internal_storage) - - runtime_info = compute_handler.get_runtime_info() + handler, internal_storage = _prepare_serverless( + name, config, backend, storage, debug, 'lithops runtime delete' + ) + runtime_info = handler.get_runtime_info() runtime_name = runtime_info['runtime_name'] - runtimes = compute_handler.list_runtimes(runtime_name) - runtimes_to_delete = [] - - for runtime in runtimes: - to_delete = True - if memory is not None and runtime[1] != int(memory): - to_delete = False - if version is not None and runtime[2] != version: - to_delete = False - if to_delete: - runtimes_to_delete.append((runtime[0], runtime[1], runtime[2])) + # Rows are (name, memory, version, worker name), see list_runtimes + runtimes_to_delete = [ + (rt[0], rt[1], rt[2]) + for rt in handler.list_runtimes(runtime_name) + if (memory is None or rt[1] == int(memory)) + and (version is None or rt[2] == version) + ] if not runtimes_to_delete: logger.info("Runtime not found") return - for runtime in runtimes_to_delete: - compute_handler.delete_runtime(runtime[0], runtime[1], runtime[2]) - runtime_key = compute_handler.get_runtime_key(runtime[0], runtime[1], runtime[2]) + for rt_name, rt_memory, rt_version in runtimes_to_delete: + handler.delete_runtime(rt_name, rt_memory, rt_version) + runtime_key = handler.get_runtime_key(rt_name, rt_memory, rt_version) internal_storage.delete_runtime_meta(runtime_key) logger.info("Runtime deleted") @@ -677,73 +904,36 @@ def delete(name, config, memory, version, backend, storage, debug): # /---------------------------------------------------------------------------/ @click.group('job') -@click.pass_context -def job(ctx): +def job(): pass -@job.command('list', context_settings=dict(ignore_unknown_options=True, allow_extra_args=True)) -@click.option('--config', '-c', default=None, help='path to yaml config file', type=click.Path(exists=True)) +@job.command('list') +@click.option( + '--config', '-c', default=None, + help='path to yaml config file', type=click.Path(exists=True) +) @click.option('--backend', '-b', default=None, help='compute backend') @click.option('--region', '-r', default=None, help='compute backend region') @click.option('--debug', '-d', is_flag=True, help='debug mode') def list_jobs(config, backend, region, debug): """ List Standalone Jobs """ - log_level = logging.INFO if not debug else logging.DEBUG - setup_lithops_logger(log_level) - - config = load_yaml_config(config) if config else None - config_ow = set_config_ow(backend=backend, region=region) - config = default_config(config_data=config, config_overwrite=config_ow, load_storage_config=False) - - if config['lithops']['mode'] != STANDALONE: - raise Exception('"lithops job list" command is only available for standalone backends. ' - f'Please use "lithops job list -b {set(STANDALONE_BACKENDS)}"') - - compute_config = extract_standalone_config(config) - compute_handler = StandaloneHandler(compute_config) - - if not compute_handler.is_initialized(): - logger.info("The backend is not initialized") - return - - compute_handler.init() - - if not compute_handler.backend.master.is_ready(): - logger.info(f"{compute_handler.backend.master} is stopped") + handler = _prepare_standalone( + config, backend, region, debug, 'lithops job list' + ) + if not _standalone_service_ready(handler): return - if not compute_handler._is_master_service_ready(): - logger.info(f"Lithops service is not running in {compute_handler.backend.master}") + logger.info(f'Listing jobs submitted to {handler.backend.master}') + job_list = handler.list_jobs() + if not job_list: + _print_table([], [], 'jobs') return - logger.info(f'Listing jobs submitted to {compute_handler.backend.master}') - job_list = compute_handler.list_jobs() - headers = job_list.pop(0) key_index = headers.index("Submitted") - - try: - import pytz - from tzlocal import get_localzone - local_tz = get_localzone() - - def convert_utc_to_local(utc_timestamp): - utc_time = datetime.strptime(utc_timestamp, '%Y-%m-%d %H:%M:%S %Z') - utc_time = utc_time.replace(tzinfo=pytz.utc) - local_time = utc_time.astimezone(local_tz) - return local_time.strftime('%Y-%m-%d %H:%M:%S %Z') - - for row in job_list: - row[key_index] = convert_utc_to_local(row[key_index]) - except ModuleNotFoundError: - pass - - sorted_data = sorted(job_list, key=lambda x: x[key_index]) - - print() - print(tabulate(sorted_data, headers=headers)) - print(f'\nTotal jobs: {len(job_list)}') + rows = _localize_and_sort_rows(job_list, key_index) + _print_table(rows, headers, 'jobs') # /---------------------------------------------------------------------------/ @@ -753,73 +943,36 @@ def convert_utc_to_local(utc_timestamp): # /---------------------------------------------------------------------------/ @click.group('worker') -@click.pass_context -def worker(ctx): +def worker(): pass -@worker.command('list', context_settings=dict(ignore_unknown_options=True, allow_extra_args=True)) -@click.option('--config', '-c', default=None, help='path to yaml config file', type=click.Path(exists=True)) +@worker.command('list') +@click.option( + '--config', '-c', default=None, + help='path to yaml config file', type=click.Path(exists=True) +) @click.option('--backend', '-b', default=None, help='compute backend') @click.option('--region', '-r', default=None, help='compute backend region') @click.option('--debug', '-d', is_flag=True, help='debug mode') def list_workers(config, backend, region, debug): - """ List Standalone Jobs """ - log_level = logging.INFO if not debug else logging.DEBUG - setup_lithops_logger(log_level) - - config = load_yaml_config(config) if config else None - config_ow = set_config_ow(backend=backend, region=region) - config = default_config(config_data=config, config_overwrite=config_ow, load_storage_config=False) - - if config['lithops']['mode'] != STANDALONE: - raise Exception('"lithops worker list" command is only available for standalone backends. ' - f'Please use "lithops worker list -b {set(STANDALONE_BACKENDS)}"') - - compute_config = extract_standalone_config(config) - compute_handler = StandaloneHandler(compute_config) - - if not compute_handler.is_initialized(): - logger.info("The backend is not initialized") - return - - compute_handler.init() - - if not compute_handler.backend.master.is_ready(): - logger.info(f"{compute_handler.backend.master} is stopped") + """ List Standalone Workers """ + handler = _prepare_standalone( + config, backend, region, debug, 'lithops worker list' + ) + if not _standalone_service_ready(handler): return - if not compute_handler._is_master_service_ready(): - logger.info(f"Lithops service is not running in {compute_handler.backend.master}") + logger.info(f'Listing available workers in {handler.backend.master}') + worker_list = handler.list_workers() + if not worker_list: + _print_table([], [], 'workers') return - logger.info(f'Listing available workers in {compute_handler.backend.master}') - worker_list = compute_handler.list_workers() - headers = worker_list.pop(0) key_index = headers.index("Created") - - try: - import pytz - from tzlocal import get_localzone - local_tz = get_localzone() - - def convert_utc_to_local(utc_timestamp): - utc_time = datetime.strptime(utc_timestamp, '%Y-%m-%d %H:%M:%S %Z') - utc_time = utc_time.replace(tzinfo=pytz.utc) - local_time = utc_time.astimezone(local_tz) - return local_time.strftime('%Y-%m-%d %H:%M:%S %Z') - - for row in worker_list: - row[key_index] = convert_utc_to_local(row[key_index]) - except ModuleNotFoundError: - pass - - sorted_data = sorted(worker_list, key=lambda x: x[key_index]) - - print() - print(tabulate(sorted_data, headers=headers)) - print(f'\nTotal workers: {len(worker_list)}') + rows = _localize_and_sort_rows(worker_list, key_index) + _print_table(rows, headers, 'workers') # /---------------------------------------------------------------------------/ @@ -829,106 +982,91 @@ def convert_utc_to_local(utc_timestamp): # /---------------------------------------------------------------------------/ @click.group('image') -@click.pass_context -def image(ctx): +def image(): pass -@image.command('build', context_settings=dict(ignore_unknown_options=True, allow_extra_args=True)) +@image.command( + 'build', + context_settings=dict(ignore_unknown_options=True, allow_extra_args=True) +) @click.argument('name', required=False) -@click.option('--file', '-f', default=None, help='file needed to build the image', type=click.Path(exists=True)) -@click.option('--config', '-c', default=None, help='path to yaml config file', type=click.Path(exists=True)) +@click.option( + '--file', '-f', default=None, + help='file needed to build the image', type=click.Path(exists=True) +) +@click.option( + '--config', '-c', default=None, + help='path to yaml config file', type=click.Path(exists=True) +) @click.option('--backend', '-b', default=None, help='compute backend') @click.option('--region', '-r', default=None, help='compute backend region') @click.option('--debug', '-d', is_flag=True, help='debug mode') -@click.option('--overwrite', '-o', is_flag=True, help='overwrite the image if it already exists') -@click.option('--include', '-i', multiple=True, help='include source:destination paths', type=str) +@click.option( + '--overwrite', '-o', is_flag=True, + help='overwrite the image if it already exists' +) +@click.option( + '--include', '-i', multiple=True, + help='include source:destination paths', type=str +) @click.pass_context def build_image(ctx, name, file, config, backend, region, debug, overwrite, include): """ build a VM image """ - setup_lithops_logger(logging.DEBUG) - if name: verify_runtime_name(name) - - config = load_yaml_config(config) if config else None - config_ow = set_config_ow(backend=backend, region=region) - config = default_config(config_data=config, config_overwrite=config_ow, load_storage_config=False) - - if config['lithops']['mode'] != STANDALONE: - raise Exception('"lithops image build" command is only available for standalone backends. ' - f'Please use "lithops image build -b {set(STANDALONE_BACKENDS)}"') + handler = _prepare_standalone( + config, backend, region, debug, 'lithops image build', + always_debug=True, + ) for src_dst_file in include: src_file, dst_file = src_dst_file.split(':') if not os.path.isfile(src_file): raise FileNotFoundError(f"The file '{src_file}' does not exist") - compute_config = extract_standalone_config(config) - compute_handler = StandaloneHandler(compute_config) - compute_handler.build_image(name, file, overwrite, include, ctx.args) - + handler.build_image(name, file, overwrite, include, ctx.args) logger.info('VM Image built') -@image.command('delete', context_settings=dict(ignore_unknown_options=True, allow_extra_args=True)) +@image.command('delete') @click.argument('name', required=True) -@click.option('--config', '-c', default=None, help='path to yaml config file', type=click.Path(exists=True)) +@click.option( + '--config', '-c', default=None, + help='path to yaml config file', type=click.Path(exists=True) +) @click.option('--backend', '-b', default=None, help='compute backend') @click.option('--region', '-r', default=None, help='compute backend region') @click.option('--debug', '-d', is_flag=True, help='debug mode') -@click.pass_context -def delete_image(ctx, name, config, backend, region, debug): +def delete_image(name, config, backend, region, debug): """ Delete a VM image """ - setup_lithops_logger(logging.DEBUG) - if name: verify_runtime_name(name) - - config = load_yaml_config(config) if config else None - config_ow = set_config_ow(backend=backend, region=region) - config = default_config(config_data=config, config_overwrite=config_ow, load_storage_config=False) - - if config['lithops']['mode'] != STANDALONE: - raise Exception('"lithops image delete" command is only available for standalone backends. ' - f'Please use "lithops image delete -b {set(STANDALONE_BACKENDS)}"') - - compute_config = extract_standalone_config(config) - compute_handler = StandaloneHandler(compute_config) - compute_handler.delete_image(name) - + handler = _prepare_standalone( + config, backend, region, debug, 'lithops image delete', + always_debug=True, + ) + handler.delete_image(name) logger.info('VM Image deleted') -@image.command('list', context_settings=dict(ignore_unknown_options=True, allow_extra_args=True)) -@click.option('--config', '-c', default=None, help='path to yaml config file', type=click.Path(exists=True)) +@image.command('list') +@click.option( + '--config', '-c', default=None, + help='path to yaml config file', type=click.Path(exists=True) +) @click.option('--backend', '-b', default=None, help='compute backend') @click.option('--region', '-r', default=None, help='compute backend region') @click.option('--debug', '-d', is_flag=True, help='debug mode') def list_images(config, backend, region, debug): """ List VM images """ - log_level = logging.INFO if not debug else logging.DEBUG - setup_lithops_logger(log_level) - - config = load_yaml_config(config) if config else None - config_ow = set_config_ow(backend=backend, region=region) - config = default_config(config_data=config, config_overwrite=config_ow, load_storage_config=False) - - if config['lithops']['mode'] != STANDALONE: - raise Exception('"lithops image build" command is only available for standalone backends. ' - f'Please use "lithops image list -b {set(STANDALONE_BACKENDS)}"') - - compute_config = extract_standalone_config(config) - compute_handler = StandaloneHandler(compute_config) - + handler = _prepare_standalone( + config, backend, region, debug, 'lithops image list' + ) logger.info('Listing all Ubuntu VM images') - vm_images = compute_handler.list_images() - + vm_images = handler.list_images() headers = ['Image Name', 'Image ID', 'Creation Date'] - - print() - print(tabulate(vm_images, headers=headers)) - print(f'\nTotal VM images: {len(vm_images)}') + _print_table(vm_images, headers, 'VM images') lithops_cli.add_command(runtime) diff --git a/lithops/serverless/backends/aws_batch/aws_batch.py b/lithops/serverless/backends/aws_batch/aws_batch.py index 14356b004..be3e06f16 100644 --- a/lithops/serverless/backends/aws_batch/aws_batch.py +++ b/lithops/serverless/backends/aws_batch/aws_batch.py @@ -544,7 +544,7 @@ def invoke(self, runtime_name, runtime_memory, payload): payload['chunksize'] = chunksize logger.debug( - f'ExecutorID {executor_id} | JobID {job_id} - Required Workers: {total_workers}' + f'{utils.log_prefix(executor_id, job_id)} - Required Workers: {total_workers}' ) job_name = '{}_{}'.format(self._format_jobdef_name(runtime_name, runtime_memory), payload['job_key']) diff --git a/lithops/serverless/backends/aws_lambda/aws_lambda.py b/lithops/serverless/backends/aws_lambda/aws_lambda.py index 05742384f..93c5ac1ba 100644 --- a/lithops/serverless/backends/aws_lambda/aws_lambda.py +++ b/lithops/serverless/backends/aws_lambda/aws_lambda.py @@ -129,13 +129,19 @@ def _create_handler_bin(remove=True): """ current_location = os.path.dirname(os.path.abspath(__file__)) main_file = os.path.join(current_location, 'entry_point.py') + logger.debug(f'Building handler zip at {os.path.abspath(LITHOPS_FUNCTION_ZIP)}') utils.create_handler_zip(LITHOPS_FUNCTION_ZIP, main_file, 'entry_point.py') with open(LITHOPS_FUNCTION_ZIP, 'rb') as action_zip: action_bin = action_zip.read() + logger.debug( + f'Handler zip loaded into memory - Size: {utils.sizeof_fmt(len(action_bin))}' + ) + if remove: os.remove(LITHOPS_FUNCTION_ZIP) + logger.debug(f'Removed temporary handler zip {LITHOPS_FUNCTION_ZIP}') return action_bin @@ -145,23 +151,31 @@ def _wait_for_function_deployed(self, func_name): Raises exception if waiting times out or if state is 'Failed' or 'Inactive' """ retries, sleep_seconds = (15, 25) if 'vpc' in self.lambda_config else (30, 5) + logger.debug( + f'Waiting for Lambda function "{func_name}" to become active' + ) while retries > 0: res = self.lambda_client.get_function(FunctionName=func_name) state = res['Configuration']['State'] if state == 'Pending': + logger.debug( + f'"{func_name}" function is being deployed... (status: {state})' + ) time.sleep(sleep_seconds) - logger.debug('"{}" function is being deployed... ' - '(status: {})'.format(func_name, res['Configuration']['State'])) retries -= 1 if retries == 0: - raise Exception('"{}" function not deployed (timed out): {}'.format(func_name, res)) - elif state == 'Failed' or state == 'Inactive': - raise Exception('"{}" function not deployed (state is "{}"): {}'.format(func_name, state, res)) + raise Exception( + f'"{func_name}" function not deployed (timed out): {res}' + ) + elif state in ('Failed', 'Inactive'): + raise Exception( + f'"{func_name}" function not deployed (state is "{state}"): {res}' + ) elif state == 'Active': break - logger.debug('Ok --> function "{}" is active'.format(func_name)) + logger.debug(f'Ok --> function "{func_name}" is active') def _get_layer(self, runtime_name): """ @@ -393,6 +407,10 @@ def _deploy_default_runtime(self, runtime_name, memory, timeout): code = self._create_handler_bin() env_vars = {t['name']: t['value'] for t in self.lambda_config['env_vars']} + logger.debug( + f'Creating Lambda function {function_name} ' + f'({utils.sizeof_fmt(len(code))} payload)' + ) try: response = self.lambda_client.create_function( FunctionName=function_name, @@ -428,16 +446,25 @@ def _deploy_default_runtime(self, runtime_name, memory, timeout): ) if response['ResponseMetadata']['HTTPStatusCode'] not in (200, 201): - raise Exception(f'An error occurred creating/updating action {runtime_name}: {response}') + raise Exception( + f'An error occurred creating/updating action {runtime_name}: {response}' + ) + logger.debug( + f'Create function request accepted for {function_name} ' + f'(HTTP {response["ResponseMetadata"]["HTTPStatusCode"]})' + ) except Exception as e: if 'ResourceConflictException' in str(e): - pass + logger.debug( + f'Lambda function {function_name} already exists, ' + 'waiting until it is active' + ) else: raise e self._wait_for_function_deployed(function_name) - logger.debug('OK --> Created lambda function {}'.format(function_name)) + logger.debug(f'OK --> Created lambda function {function_name}') def _deploy_container_runtime(self, runtime_name, memory, timeout): """ diff --git a/lithops/serverless/backends/azure_containers/azure_containers.py b/lithops/serverless/backends/azure_containers/azure_containers.py index ce3bb7dc9..f7a850db2 100644 --- a/lithops/serverless/backends/azure_containers/azure_containers.py +++ b/lithops/serverless/backends/azure_containers/azure_containers.py @@ -70,26 +70,35 @@ def _check_az_cli(self): def _run_az_command(self, cmd, return_json=False, return_result=False): """ - Run an Azure CLI command using shell=True. + Run an Azure CLI command. Uses subprocess.run so az progress on + stderr cannot fill a pipe and deadlock (unlike check_call + PIPE). """ self._check_az_cli() - quiet = logger.getEffectiveLevel() != logging.DEBUG - kwargs = {'shell': True, 'encoding': 'UTF-8', 'stderr': sp.PIPE} - if quiet and not (return_json or return_result): - kwargs['stdout'] = sp.DEVNULL + debug = logger.getEffectiveLevel() == logging.DEBUG + capture = return_json or return_result or not debug + + logger.debug(f'Running Azure CLI: {cmd}') try: - if return_json or return_result: - result = sp.check_output(cmd, **kwargs) - else: - sp.check_call(cmd, **kwargs) - return None + completed = sp.run( + cmd, + shell=True, + encoding='UTF-8', + stdout=sp.PIPE if capture else None, + stderr=sp.PIPE if capture else None, + stdin=sp.DEVNULL, + check=True, + ) except sp.CalledProcessError as e: err_msg = f'Azure CLI command failed: {cmd}' - if e.stderr: - err_msg += f'\n{e.stderr.strip()}' + detail = ((e.stderr or e.stdout) or '').strip() + if detail: + err_msg += f'\n{detail}' raise Exception(err_msg) from e - result = result.strip() + if not (return_json or return_result): + return None + + result = (completed.stdout or '').strip() if return_json: try: return json.loads(result) @@ -97,9 +106,7 @@ def _run_az_command(self, cmd, return_json=False, return_result=False): raise Exception( f'Failed to parse Azure CLI output as JSON: {result}' ) from e - if return_result: - return result.replace('"', '') - return result + return result.replace('"', '') def _format_containerapp_name(self, runtime_name, runtime_memory, version=__version__): """ @@ -120,19 +127,39 @@ def _get_managed_environment_id(self): f'--query id --only-show-errors') return self._run_az_command(cmd, return_result=True) - def _containerapp_exists(self, containerapp_name): + def _get_containerapp_provisioning_state(self, containerapp_name): cmd = (f'az containerapp show --name {containerapp_name} ' - f'--resource-group {self.resource_group} --only-show-errors') - kwargs = {'shell': True} - if logger.getEffectiveLevel() != logging.DEBUG: - kwargs['stderr'] = sp.DEVNULL - kwargs['stdout'] = sp.DEVNULL + f'--resource-group {self.resource_group} ' + f'--query properties.provisioningState -o tsv') try: - sp.check_call(cmd, **kwargs) - return True - except sp.CalledProcessError: - logger.debug(f'Container app {containerapp_name} not found, will create it') - return False + state = self._run_az_command(cmd, return_result=True) + return (state or '').strip() or None + except Exception: + return None + + def _wait_until_containerapp_idle(self, containerapp_name, timeout=360): + """ + Wait until the app is not creating, updating, or deleting. + Returns the provisioning state, or None if the app does not exist. + """ + deadline = time.time() + timeout + busy = ('inprogress', 'deleting') + logged = False + while time.time() < deadline: + state = self._get_containerapp_provisioning_state(containerapp_name) + if not state or state.lower() not in busy: + return state + if not logged: + logger.info( + f'Container app {containerapp_name} has a provisioning ' + f'operation in progress, waiting for it to finish' + ) + logged = True + logger.debug( + f'Container app {containerapp_name} provisioning state: {state}' + ) + time.sleep(15) + return self._get_containerapp_provisioning_state(containerapp_name) def _get_default_runtime_image_name(self): """ @@ -274,26 +301,44 @@ def _create_app(self, runtime_name, memory, timeout): deployed = False retries = 0 last_error = None + cmd = None while retries < 10: try: - if self._containerapp_exists(containerapp_name): + state = self._wait_until_containerapp_idle(containerapp_name) + if state and state.lower() == 'failed': + logger.warning( + f'Container app {containerapp_name} is in Failed state, recreating' + ) + del_cmd = ( + f'az containerapp delete --name {containerapp_name} ' + f'--resource-group {self.resource_group} -y -o none' + ) + self._run_az_command(del_cmd) + self._wait_until_containerapp_idle(containerapp_name) + state = None + + if state: logger.debug(f'Container app {containerapp_name} already exists, updating') cmd = (f'az containerapp update --name {containerapp_name} ' f'--resource-group {self.resource_group} ' - f'--yaml {config.CA_JSON_LOCATION} --only-show-errors') + f'--yaml {config.CA_JSON_LOCATION} -o none') else: cmd = (f'az containerapp create --name {containerapp_name} ' f'--resource-group {self.resource_group} ' - f'--yaml {config.CA_JSON_LOCATION} --only-show-errors') + f'--yaml {config.CA_JSON_LOCATION} -o none') self._run_az_command(cmd) os.remove(config.CA_JSON_LOCATION) deployed = True break except Exception as e: last_error = e + in_progress = 'ContainerAppOperationInProgress' in str(e) logger.warning(f'Container app deploy attempt {retries + 1} failed: {e}') - time.sleep(10) + if in_progress: + self._wait_until_containerapp_idle(containerapp_name) + else: + time.sleep(15) retries += 1 if not deployed: @@ -303,9 +348,11 @@ def delete_runtime(self, runtime_name, memory, version=__version__): """ Deletes a runtime """ - logger.info(f'Deleting runtime: {runtime_name} - {memory}MB') + logger.info(f'Deleting runtime: {runtime_name} - {memory}MB (this may take several minutes)') containerapp_name = self._format_containerapp_name(runtime_name, memory, version) - cmd = f'az containerapp delete --name {containerapp_name} --resource-group {self.resource_group} -y --only-show-errors' + self._wait_until_containerapp_idle(containerapp_name) + cmd = (f'az containerapp delete --name {containerapp_name} ' + f'--resource-group {self.resource_group} -y -o none') self._run_az_command(cmd) try: diff --git a/lithops/serverless/backends/azure_functions/azure_functions.py b/lithops/serverless/backends/azure_functions/azure_functions.py index 8b23a8867..532865237 100644 --- a/lithops/serverless/backends/azure_functions/azure_functions.py +++ b/lithops/serverless/backends/azure_functions/azure_functions.py @@ -73,26 +73,35 @@ def _check_az_cli(self): def _run_az_command(self, cmd, return_json=False, return_result=False): """ - Run an Azure CLI command using shell=True. + Run an Azure CLI command. Uses subprocess.run so az progress on + stderr cannot fill a pipe and deadlock (unlike check_call + PIPE). """ self._check_az_cli() - quiet = logger.getEffectiveLevel() != logging.DEBUG - kwargs = {'shell': True, 'encoding': 'UTF-8', 'stderr': sp.PIPE} - if quiet and not (return_json or return_result): - kwargs['stdout'] = sp.DEVNULL + debug = logger.getEffectiveLevel() == logging.DEBUG + capture = return_json or return_result or not debug + + logger.debug(f'Running Azure CLI: {cmd}') try: - if return_json or return_result: - result = sp.check_output(cmd, **kwargs) - else: - sp.check_call(cmd, **kwargs) - return None + completed = sp.run( + cmd, + shell=True, + encoding='UTF-8', + stdout=sp.PIPE if capture else None, + stderr=sp.PIPE if capture else None, + stdin=sp.DEVNULL, + check=True, + ) except sp.CalledProcessError as e: err_msg = f'Azure CLI command failed: {cmd}' - if e.stderr: - err_msg += f'\n{e.stderr.strip()}' + detail = ((e.stderr or e.stdout) or '').strip() + if detail: + err_msg += f'\n{detail}' raise Exception(err_msg) from e - result = result.strip() + if not (return_json or return_result): + return None + + result = (completed.stdout or '').strip() if return_json: try: return json.loads(result) @@ -100,19 +109,19 @@ def _run_az_command(self, cmd, return_json=False, return_result=False): raise Exception( f'Failed to parse Azure CLI output as JSON: {result}' ) from e - if return_result: - return result.replace('"', '') - return result + return result.replace('"', '') def _function_app_exists(self, function_name): cmd = (f'az functionapp show --name {function_name} ' f'--resource-group {self.resource_group}') - kwargs = {} - if logger.getEffectiveLevel() != logging.DEBUG: - kwargs['stderr'] = sp.DEVNULL - kwargs['stdout'] = sp.DEVNULL try: - sp.check_call(cmd, shell=True, **kwargs) + # Probe only; hide az CLI's ResourceNotFound when the app is missing. + sp.check_call( + cmd, + shell=True, + stdout=sp.DEVNULL, + stderr=sp.DEVNULL, + ) return True except sp.CalledProcessError: logger.debug(f'Function app {function_name} not found, will create it') diff --git a/lithops/serverless/backends/code_engine/code_engine.py b/lithops/serverless/backends/code_engine/code_engine.py index 2fbae6a48..02df47a23 100644 --- a/lithops/serverless/backends/code_engine/code_engine.py +++ b/lithops/serverless/backends/code_engine/code_engine.py @@ -823,7 +823,7 @@ def invoke(self, docker_image_name, runtime_memory, job_payload): job_payload['chunksize'] = chunksize logger.debug( - f'ExecutorID {executor_id} | JobID {job_id} - Required Workers: {total_workers}' + f'{utils.log_prefix(executor_id, job_id)} - Required Workers: {total_workers}' ) jobdef_name = self._format_jobdef_name(docker_image_name, runtime_memory) diff --git a/lithops/serverless/backends/gcp_cloudrun/cloudrun.py b/lithops/serverless/backends/gcp_cloudrun/cloudrun.py index 2a5ec56b6..d8d9667bb 100644 --- a/lithops/serverless/backends/gcp_cloudrun/cloudrun.py +++ b/lithops/serverless/backends/gcp_cloudrun/cloudrun.py @@ -316,9 +316,9 @@ def invoke(self, runtime_name, runtime_memory, payload, return_result=False): service_url, id_token = self._get_url_and_token(service_name) if exec_id and job_id and call_id: - logger.debug(f'ExecutorID {exec_id} | JobID {job_id} - Invoking function call {call_id}') + logger.debug(f'{utils.log_prefix(exec_id, job_id)} - Invoking function call {call_id}') elif exec_id and job_id: - logger.debug(f'ExecutorID {exec_id} | JobID {job_id} - Invoking function') + logger.debug(f'{utils.log_prefix(exec_id, job_id)} - Invoking function') else: logger.debug('Invoking function') diff --git a/lithops/serverless/backends/gcp_functions/gcp_functions.py b/lithops/serverless/backends/gcp_functions/gcp_functions.py index bf531ca2b..564e5877c 100644 --- a/lithops/serverless/backends/gcp_functions/gcp_functions.py +++ b/lithops/serverless/backends/gcp_functions/gcp_functions.py @@ -432,9 +432,9 @@ def invoke(self, runtime_name, runtime_memory, payload={}, return_result=False): job_id = payload.get('job_id') if exec_id and job_id and call_id: - logger.debug(f'ExecutorID {exec_id} | JobID {job_id} - Invoking function call {call_id}') + logger.debug(f'{utils.log_prefix(exec_id, job_id)} - Invoking function call {call_id}') elif exec_id and job_id: - logger.debug(f'ExecutorID {exec_id} | JobID {job_id} - Invoking function') + logger.debug(f'{utils.log_prefix(exec_id, job_id)} - Invoking function') else: logger.debug('Invoking function') diff --git a/lithops/serverless/backends/k8s/k8s.py b/lithops/serverless/backends/k8s/k8s.py index b81dd3555..6867b55c2 100644 --- a/lithops/serverless/backends/k8s/k8s.py +++ b/lithops/serverless/backends/k8s/k8s.py @@ -684,7 +684,7 @@ def invoke(self, docker_image_name, runtime_memory, job_payload): total_workers = min(max_workers, total_calls // chunksize + (total_calls % chunksize > 0)) logger.debug( - f'ExecutorID {executor_id} | JobID {job_id} - Required Workers: {total_workers}' + f'{utils.log_prefix(executor_id, job_id)} - Required Workers: {total_workers}' ) activation_id = f'lithops-{job_key.lower()}' @@ -714,8 +714,10 @@ def invoke(self, docker_image_name, runtime_memory, job_payload): container['resources']['limits']['memory'] = f'{runtime_memory}Mi' container['resources']['limits']['cpu'] = str(self.k8s_config['runtime_cpu']) - logger.debug(f'ExecutorID {executor_id} | JobID {job_id} - Going ' - f'to run {total_calls} activations in {total_workers} workers') + logger.debug( + f'{utils.log_prefix(executor_id, job_id)} - Going ' + f'to run {total_calls} activations in {total_workers} workers' + ) if not all(key in self.k8s_config for key in ["docker_user", "docker_password"]): del job_res['spec']['template']['spec']['imagePullSecrets'] diff --git a/lithops/serverless/backends/knative/knative.py b/lithops/serverless/backends/knative/knative.py index 7621f7f23..02eb5ab28 100644 --- a/lithops/serverless/backends/knative/knative.py +++ b/lithops/serverless/backends/knative/knative.py @@ -678,11 +678,12 @@ def invoke(self, runtime_name, memory, payload, return_result=False): conn = http.client.HTTPConnection(parsed_url.netloc) if exec_id and job_id and call_ids: - logger.debug('ExecutorID {} | JobID {} - Invoking function call {}' - .format(exec_id, job_id, ', '.join(call_ids))) + logger.debug( + f'{utils.log_prefix(exec_id, job_id)} - Invoking function call ' + f'{", ".join(call_ids)}' + ) elif exec_id and job_id: - logger.debug('ExecutorID {} | JobID {} - Invoking function' - .format(exec_id, job_id)) + logger.debug(f'{utils.log_prefix(exec_id, job_id)} - Invoking function') else: logger.debug('Invoking function') @@ -704,8 +705,10 @@ def invoke(self, runtime_name, memory, payload, return_result=False): elif resp_status == 404: raise Exception("Lithops runtime is not deployed in your k8s cluster") else: - logger.debug('ExecutorID {} | JobID {} - Function call {} failed ({}). Retrying request' - .format(exec_id, job_id, ', '.join(call_ids), resp_status)) + logger.debug( + f'{utils.log_prefix(exec_id, job_id)} - Function call {", ".join(call_ids)} ' + f'failed ({resp_status}). Retrying request' + ) def get_runtime_key(self, runtime_name, runtime_memory, version=__version__): """ diff --git a/lithops/serverless/serverless.py b/lithops/serverless/serverless.py index 4f270cd47..8695d7233 100644 --- a/lithops/serverless/serverless.py +++ b/lithops/serverless/serverless.py @@ -17,114 +17,106 @@ import logging import importlib +from typing import Any, Dict logger = logging.getLogger(__name__) class ServerlessHandler: """ - A ServerlessHandler object is used by invokers and other components to access - underlying serverless backend without exposing the implementation details. + A ServerlessHandler object is used by invokers and other components to + access the underlying serverless backend without exposing implementation + details. """ - def __init__(self, servereless_config, internal_storage): - self.config = servereless_config + def __init__(self, serverless_config: Dict[str, Any], internal_storage): + self.config = serverless_config self.backend_name = self.config['backend'] - self.backend = None + self.backend = self._load_backend(internal_storage) + def _load_backend(self, internal_storage): + """Builds the backend the configuration asks for""" try: module_location = f'lithops.serverless.backends.{self.backend_name}' sb_module = importlib.import_module(module_location) - ServerlessBackend = getattr(sb_module, 'ServerlessBackend') - self.backend = ServerlessBackend(self.config[self.backend_name], internal_storage) - - except Exception as e: - logger.error("There was an error trying to create the {} " - "serverless backend".format(self.backend_name)) - raise e + serverless_backend_cls = getattr(sb_module, 'ServerlessBackend') + return serverless_backend_cls( + self.config[self.backend_name], internal_storage + ) + except Exception: + logger.error( + f"There was an error trying to create the {self.backend_name} " + "serverless backend", + exc_info=True, + ) + raise + + def _call_backend(self, method: str, *args): + """ + Calls an optional backend method, so that a backend only implements + the hooks it needs + """ + fn = getattr(self.backend, method, None) + if fn is not None: + return fn(*args) + return None def init(self): - """ - Init tasks for serverless batch backends - """ - pass + """Nothing to initialize: serverless backends are ready to invoke""" def pre_invoke(self, job): - """ - Pre-invocation task executed just before the actual parallel invocation - in the serverless FaaS backends. - """ - runtime_name = job.runtime_name - runtime_memory = job.runtime_memory - - if hasattr(self.backend, 'pre_invoke'): - self.backend.pre_invoke(runtime_name, runtime_memory) - - def invoke(self, job_payload): - """ - Invoke -- return information about this invocation - """ - runtime_name = job_payload['runtime_name'] - runtime_memory = job_payload['runtime_memory'] - - return self.backend.invoke(runtime_name, runtime_memory, job_payload) - - def build_runtime(self, runtime_name, file, extra_args=[]): - """ - Wrapper method to build a new runtime for the compute backend. - return: the name of the runtime - """ - self.backend.build_runtime(runtime_name, file, extra_args) - - def deploy_runtime(self, runtime_name, memory, timeout): - """ - Wrapper method to deploy a runtime in the compute backend. - return: the name of the runtime - """ - return self.backend.deploy_runtime(runtime_name, memory, timeout=timeout) - - def delete_runtime(self, runtime_name, memory, version): - """ - Wrapper method to delete a runtime in the compute backend - """ + """Runs the pre-invoke hook of the backend, if it has one""" + self._call_backend('pre_invoke', job.runtime_name, job.runtime_memory) + + def invoke(self, job_payload: Dict[str, Any]): + """Invokes a job, and returns the activation id of the invocation""" + return self.backend.invoke( + job_payload['runtime_name'], + job_payload['runtime_memory'], + job_payload, + ) + + def build_runtime(self, runtime_name: str, file: str, extra_args=None): + """Builds the runtime image the jobs will run in""" + self.backend.build_runtime(runtime_name, file, extra_args or []) + + def deploy_runtime(self, runtime_name: str, memory: int, timeout: int): + """Deploys a runtime and returns its metadata""" + return self.backend.deploy_runtime( + runtime_name, memory, timeout=timeout + ) + + def delete_runtime(self, runtime_name: str, memory: int, version: str): + """Deletes a deployed runtime""" self.backend.delete_runtime(runtime_name, memory, version) def clean(self, **kwargs): - """ - Wrapper method to clean the compute backend - """ + """Deletes every runtime and every resource the backend created""" self.backend.clean(**kwargs) def clear(self, job_keys=None, exception=None): """ - Wrapper method to clear the compute backend + Releases the backend resources of the given jobs, for the backends + that hold any """ - if hasattr(self.backend, 'clear'): - self.backend.clear(job_keys) + self._call_backend('clear', job_keys) - def list_runtimes(self, runtime_name='all'): - """ - Wrapper method to list deployed runtime in the compute backend - """ + def list_runtimes(self, runtime_name: str = 'all'): + """Lists the runtimes deployed in the backend""" return self.backend.list_runtimes(runtime_name) - def get_runtime_key(self, runtime_name, memory, version): + def get_runtime_key(self, runtime_name: str, memory: int, version: str): """ - Wrapper method that returns a formated string that represents the runtime key. - Each backend has its own runtime key format. Used to store runtime metadata - into the storage + Returns a formatted string that represents the runtime key. + Each backend has its own runtime key format. Used to store + runtime metadata in storage. """ return self.backend.get_runtime_key(runtime_name, memory, version) def get_runtime_info(self): - """ - Wrapper method that returns a dictionary with all the runtime information - set in config - """ + """Returns the runtime limits the executor reports to the user""" return self.backend.get_runtime_info() def get_backend_type(self): - """ - Wrapper method that returns the type of the backend (Batch or FaaS) - """ + """Returns whether the backend is invoked per call or with a job""" return self.backend.type diff --git a/lithops/standalone/__init__.py b/lithops/standalone/__init__.py index c0f4f6d71..1bcdec167 100644 --- a/lithops/standalone/__init__.py +++ b/lithops/standalone/__init__.py @@ -1,4 +1,4 @@ from .standalone import StandaloneHandler from .utils import LithopsValidationError -__all__ = ['StandaloneHandler', LithopsValidationError] +__all__ = ['StandaloneHandler', 'LithopsValidationError'] diff --git a/lithops/standalone/keeper.py b/lithops/standalone/keeper.py index beb4495b6..bd6008f3b 100644 --- a/lithops/standalone/keeper.py +++ b/lithops/standalone/keeper.py @@ -19,6 +19,8 @@ import time import threading import logging +from typing import Any, Callable, Dict, Optional + from lithops.standalone import StandaloneHandler from lithops.constants import JOBS_DIR from lithops.standalone.utils import JobStatus @@ -29,10 +31,19 @@ class BudgetKeeper(threading.Thread): """ - BudgetKeeper class used to automatically stop the VM instance + Background thread that stops the VM instance it runs on once it has been + idle for long enough, so that a forgotten or misconfigured run does not + keep paying for it """ - def __init__(self, config, instance_data, stop_callback=None, delete_callback=None): - threading.Thread.__init__(self) + + def __init__( + self, + config: Dict[str, Any], + instance_data: Dict[str, Any], + stop_callback: Optional[Callable] = None, + delete_callback: Optional[Callable] = None, + ): + super().__init__() self.last_usage_time = time.time() self.standalone_config = config @@ -43,85 +54,126 @@ def __init__(self, config, instance_data, stop_callback=None, delete_callback=No self.hard_dismantle_timeout = config['hard_dismantle_timeout'] self.exec_mode = config['exec_mode'] - self.runing = False + self.running = False self.jobs = {} self.time_to_dismantle = self.hard_dismantle_timeout self.standalone_handler = StandaloneHandler(self.standalone_config) self.instance = self.standalone_handler.backend.get_instance(**instance_data) - logger.debug(f"Starting BudgetKeeper for {self.instance.name} ({self.instance.private_ip}), " - f"instance ID: {self.instance.instance_id}") - logger.debug(f"Delete {self.instance.name} on dismantle: {self.instance.delete_on_dismantle}") + logger.debug( + f"Starting BudgetKeeper for {self.instance.name} " + f"({self.instance.private_ip}), instance ID: {self.instance.instance_id}" + ) + logger.debug( + f"Delete {self.instance.name} on dismantle: " + f"{self.instance.delete_on_dismantle}" + ) def get_time_to_dismantle(self): + """Returns the seconds left before the instance is stopped""" return self.time_to_dismantle def add_job(self, job_key): + """Marks a job as running, which pushes the countdown forward""" self.last_usage_time = time.time() self.jobs[job_key] = JobStatus.RUNNING.value def set_job_done(self, job_key): + """Marks a job as done, which starts the idle countdown""" self.last_usage_time = time.time() self.jobs[job_key] = JobStatus.DONE.value - def run(self): - self.runing = True - jobs_running = False - - logger.debug("BudgetKeeper started") - + def _all_jobs_done(self): + """True when there has been at least one job and none is running""" + return bool(self.jobs) and all( + status == JobStatus.DONE.value for status in self.jobs.values() + ) + + def _mark_finished_jobs(self): + """ + Marks as done the jobs whose runner left a done file behind. Iterates + over a snapshot, because the service adds jobs from its own threads + """ + for job_key in list(self.jobs.keys()): + done_file = os.path.join(JOBS_DIR, job_key + '.done') + if os.path.isfile(done_file): + self.jobs[job_key] = JobStatus.DONE.value + + def _log_dismantle_settings(self): + """Reports the timeouts the countdown will use""" if self.auto_dismantle: - logger.debug('Auto dismantle activated - Soft timeout: {}s, Hard Timeout: {}s' - .format(self.soft_dismantle_timeout, self.hard_dismantle_timeout)) + logger.debug( + f'Auto dismantle activated - Soft timeout: ' + f'{self.soft_dismantle_timeout}s, Hard Timeout: ' + f'{self.hard_dismantle_timeout}s' + ) else: # If auto_dismantle is deactivated, the VM will be always automatically # stopped after hard_dismantle_timeout. This will prevent the VM # being started forever due a wrong configuration - logger.debug(f'Auto dismantle deactivated - Hard Timeout: {self.hard_dismantle_timeout}s') + logger.debug( + f'Auto dismantle deactivated - ' + f'Hard Timeout: {self.hard_dismantle_timeout}s' + ) - while self.runing: - time_since_last_usage = time.time() - self.last_usage_time + def run(self): + """ + Counts down to the moment the instance is stopped, for as long as + nothing pushes the countdown forward + """ + self.running = True + jobs_running = False - for job_key in self.jobs.keys(): - done = os.path.join(JOBS_DIR, job_key + '.done') - if os.path.isfile(done): - self.jobs[job_key] = JobStatus.DONE.value + logger.debug("BudgetKeeper started") + self._log_dismantle_settings() - if len(self.jobs) > 0 and all(value == JobStatus.DONE.value for value in self.jobs.values()) \ - and self.auto_dismantle: + while self.running: + time_since_last_usage = time.time() - self.last_usage_time + self._mark_finished_jobs() - # here we need to catch a moment when number of running JOBS become zero. - # when it happens we reset countdown back to soft_dismantle_timeout + if self._all_jobs_done() and self.auto_dismantle: + # Catch the moment when the number of running jobs becomes zero + # and reset the countdown back to soft_dismantle_timeout. if jobs_running: jobs_running = False self.last_usage_time = time.time() time_since_last_usage = time.time() - self.last_usage_time - - self.time_to_dismantle = int(self.soft_dismantle_timeout - time_since_last_usage) + self.time_to_dismantle = int( + self.soft_dismantle_timeout - time_since_last_usage + ) else: - self.time_to_dismantle = int(self.hard_dismantle_timeout - time_since_last_usage) + self.time_to_dismantle = int( + self.hard_dismantle_timeout - time_since_last_usage + ) jobs_running = True if self.time_to_dismantle > 0: - logger.debug(f"Time to dismantle: {self.time_to_dismantle} seconds") + logger.debug( + f"Time to dismantle: {self.time_to_dismantle} seconds" + ) check_interval = min(60, max(self.time_to_dismantle / 10, 1)) time.sleep(check_interval) else: self.stop_instance() def stop_instance(self): + """ + Stops or deletes the instance, telling whoever registered a callback + first so that it can report the instance is going away + """ logger.debug("Dismantling setup") if self.instance.delete_on_dismantle: - self.delete_callback() if self.delete_callback is not None else None - else: - self.stop_callback() if self.stop_callback is not None else None + if self.delete_callback is not None: + self.delete_callback() + elif self.stop_callback is not None: + self.stop_callback() try: self.instance.stop() - self.runing = False + self.running = False except Exception as e: logger.debug(f"Dismantle error {e}") time.sleep(5) diff --git a/lithops/standalone/master.py b/lithops/standalone/master.py index c25d7736a..4a80176ea 100644 --- a/lithops/standalone/master.py +++ b/lithops/standalone/master.py @@ -37,7 +37,6 @@ from lithops.standalone.keeper import BudgetKeeper from lithops.config import extract_standalone_config from lithops.standalone.standalone import StandaloneHandler -from lithops.version import __version__ as lithops_version from lithops.constants import ( CPU_COUNT, LITHOPS_TEMP_DIR, @@ -46,11 +45,12 @@ SA_MASTER_SERVICE_PORT, SA_WORKER_SERVICE_PORT, SA_CONFIG_FILE, - SA_MASTER_DATA_FILE + SA_MASTER_DATA_FILE, ) from lithops.utils import ( verify_runtime_name, - setup_lithops_logger + setup_lithops_logger, + log_prefix, ) from lithops.standalone.utils import ( JobStatus, @@ -61,51 +61,104 @@ install_script_kwargs_from_config, ) -os.makedirs(LITHOPS_TEMP_DIR, exist_ok=True) - -log_format = "%(asctime)s\t[%(levelname)s] %(name)s:%(lineno)s -- %(message)s" -setup_lithops_logger(logging.DEBUG, filename=SA_MASTER_LOG_FILE, log_format=log_format) logger = logging.getLogger('lithops.standalone.master') app = flask.Flask(__name__) MAX_INSTANCE_CREATE_RETRIES = 2 JOB_MONITOR_CHECK_INTERVAL = 1 +_LOG_FORMAT = ( + "%(asctime)s\t[%(levelname)s] %(name)s:%(lineno)s -- %(message)s" +) +_NOT_A_DICT = 'The action did not receive a dictionary as an argument.' redis_client = None budget_keeper = None master_ip = None +def _configure_logging(): + """Sends everything this service logs to the master log""" + os.makedirs(LITHOPS_TEMP_DIR, exist_ok=True) + setup_lithops_logger( + logging.DEBUG, filename=SA_MASTER_LOG_FILE, log_format=_LOG_FORMAT + ) + + +def _json_body(): + """Returns the body of the request, or None when it is not JSON""" + return flask.request.get_json(force=True, silent=True) + + +def _require_dict(payload): + """Returns an error response when the body is not a dictionary""" + if not isinstance(payload, dict): + return error(_NOT_A_DICT) + return None + + +def _map_if_any(fn, items): + """ + Runs fn over every item in parallel, reporting the ones that failed. + Nothing here is fatal: a worker that cannot be reached is a row missing + from a listing or a stop request nobody answered, not a failed call + """ + if not items: + return + + with ThreadPoolExecutor(len(items)) as executor: + futures = [executor.submit(fn, item) for item in items] + + # Results are only read once every thread is done, since a lazy map() + # would drop the errors on the floor + for item, future in zip(items, futures): + try: + future.result() + except Exception as e: + logger.error(f'Could not process {item}: {e}') + + +def _worker_key(worker_name): + """Returns the redis key holding the state of a worker""" + return f"worker:{worker_name}" + + +def _job_key_id(job_key): + """Returns the redis key holding the state of a job""" + return f"job:{job_key}" + + # /---------------------------------------------------------------------------/ # Workers # /---------------------------------------------------------------------------/ def is_worker_free(worker_private_ip): """ - Checks if the Lithops service is ready and free in the worker VM instance + True when the Lithops service of a worker answers and has a free process. + A worker that cannot be reached counts as not free, as there is no way to + give it any work """ url = f"http://{worker_private_ip}:{SA_WORKER_SERVICE_PORT}/ping" try: - r = requests.get(url, timeout=0.5) - resp = r.json() + resp = requests.get(url, timeout=0.5).json() logger.debug(f'Worker processes status from {worker_private_ip}: {resp}') - return True if resp.get('free', 0) > 0 else False - except Exception: + return resp.get('free', 0) > 0 + except Exception as e: + logger.debug(f'Worker {worker_private_ip} did not answer: {e}') return False def get_worker_ttd(worker_private_ip): """ - Checks if the Lithops service is ready and free in the worker VM instance + Returns the seconds left before a worker stops itself, asking the worker + unless it is this very instance, and Unknown when it cannot be asked """ try: if master_ip == worker_private_ip: ttd = str(budget_keeper.get_time_to_dismantle()) else: url = f"http://{worker_private_ip}:{SA_WORKER_SERVICE_PORT}/ttd" - r = requests.get(url, timeout=0.5) - ttd = r.text + ttd = requests.get(url, timeout=0.5).text logger.debug(f'Worker TTD from {worker_private_ip}: {ttd}') return ttd except Exception as e: @@ -115,35 +168,35 @@ def get_worker_ttd(worker_private_ip): @app.route('/worker/list', methods=['GET']) def list_workers(): - """ - Returns the current workers list - """ + """Returns a table of every worker the master knows about""" logger.debug('Listing workers') - budget_keeper.last_usage_time = time.time() - result = [['Worker Name', 'Created', 'Instance Type', 'Processes', 'Runtime', 'Mode', 'Status', 'TTD']] + result = [[ + 'Worker Name', 'Created', 'Instance Type', 'Processes', + 'Runtime', 'Mode', 'Status', 'TTD', + ]] def get_worker(worker): + """Appends the row of one worker to the table""" worker_data = redis_client.hgetall(worker) - name = worker_data['name'] - status = worker_data['status'] - private_ip = worker_data['private_ip'] - ttd = get_worker_ttd(private_ip) + ttd = get_worker_ttd(worker_data['private_ip']) ttd = ttd if ttd in ["Unknown", "Disabled"] else ttd + "s" - timestamp = float(worker_data['created']) - created = datetime.fromtimestamp(timestamp).strftime('%Y-%m-%d %H:%M:%S UTC') - instance_type = worker_data['instance_type'] - worker_processes = str(worker_data['worker_processes']) - exec_mode = worker_data['exec_mode'] - runtime = worker_data['runtime'] - result.append((name, created, instance_type, worker_processes, runtime, exec_mode, status, ttd)) - - workers = redis_client.keys('worker:*') - if workers: - with ThreadPoolExecutor(len(workers)) as ex: - ex.map(get_worker, workers) - + created = datetime.fromtimestamp( + float(worker_data['created']) + ).strftime('%Y-%m-%d %H:%M:%S UTC') + result.append(( + worker_data['name'], + created, + worker_data['instance_type'], + str(worker_data['worker_processes']), + worker_data['runtime'], + worker_data['exec_mode'], + worker_data['status'], + ttd, + )) + + _map_if_any(get_worker, redis_client.keys('worker:*')) logger.debug(f"workers: {result}") return flask.jsonify(result) @@ -151,28 +204,31 @@ def get_worker(worker): @app.route('/worker/get', methods=['GET']) def get_workers(): """ - Returns the number of free workers + Returns the workers that are free and of the shape the caller asked for, + which is how reuse mode finds the workers a new job can run on """ budget_keeper.last_usage_time = time.time() workers = redis_client.keys('worker:*') logger.debug(f'Getting workers - Total workers: {len(workers)}') - payload = flask.request.get_json(force=True, silent=True) - if payload and not isinstance(payload, dict): - return error('The action did not receive a dictionary as an argument.') + payload = _json_body() + bad_request = _require_dict(payload) + if bad_request is not None: + return bad_request worker_instance_type = payload['worker_instance_type'] worker_processes = payload['worker_processes'] runtime_name = payload['runtime_name'] active_workers = [] - for worker in workers: worker_data = redis_client.hgetall(worker) - if worker_data['instance_type'] == worker_instance_type \ - and worker_data['runtime'] == runtime_name \ - and int(worker_data['worker_processes']) == int(worker_processes): + if ( + worker_data['instance_type'] == worker_instance_type + and worker_data['runtime'] == runtime_name + and int(worker_data['worker_processes']) == int(worker_processes) + ): active_workers.append(worker_data) worker_type = f'{worker_instance_type}-{worker_processes}-{runtime_name}' @@ -181,33 +237,29 @@ def get_workers(): free_workers = [] def check_worker(worker_data): + """Keeps a worker that still has a free process""" if is_worker_free(worker_data['private_ip']): - free_workers.append( - ( - worker_data['name'], - worker_data['private_ip'], - worker_data['instance_id'], - worker_data['ssh_credentials'], - worker_data['instance_type'], - runtime_name - ) - ) - - if active_workers: - with ThreadPoolExecutor(len(active_workers)) as ex: - ex.map(check_worker, active_workers) - + free_workers.append(( + worker_data['name'], + worker_data['private_ip'], + worker_data['instance_id'], + worker_data['ssh_credentials'], + worker_data['instance_type'], + runtime_name, + )) + + _map_if_any(check_worker, active_workers) logger.debug(f'Free workers for {worker_type}: {len(free_workers)}') response = flask.jsonify(free_workers) response.status_code = 200 - return response def _redis_field(value): """ - Redis hash values must be bytes, str, int, or float. + Returns a config value as something a redis hash accepts, which is only + bytes, str, int or float """ if isinstance(value, (dict, list)): return json.dumps(value) @@ -218,16 +270,20 @@ def _redis_field(value): def save_worker(worker, standalone_config, work_queue_name): """ - Saves the worker instance with the provided data in redis + Registers a worker in redis, which is where every listing and every + lookup of a free worker reads from. The backend section is left out, as + it holds the credentials of the account """ config = copy.deepcopy(standalone_config) del config[config['backend']] config = {key: _redis_field(value) for key, value in config.items()} - worker_processes = CPU_COUNT if worker.config['worker_processes'] == 'AUTO' \ + worker_processes = ( + CPU_COUNT if worker.config['worker_processes'] == 'AUTO' else worker.config['worker_processes'] + ) - redis_client.hset(f"worker:{worker.name}", mapping={ + redis_client.hset(_worker_key(worker.name), mapping={ 'name': worker.name, 'status': WorkerStatus.STARTING.value, 'private_ip': worker.private_ip or '', @@ -237,37 +293,74 @@ def save_worker(worker, standalone_config, work_queue_name): 'created': str(time.time()), 'ssh_credentials': json.dumps(worker.ssh_credentials), 'queue_name': work_queue_name, - 'err': "", **config, + 'err': "", + **config, }) +def _worker_vm_data(instance, work_queue_name): + """Returns the data a worker VM needs to reach the master and its queue""" + return { + 'name': instance.name, + 'private_ip': instance.private_ip, + 'instance_id': instance.instance_id, + 'ssh_credentials': instance.ssh_credentials, + 'instance_type': instance.instance_type, + 'master_ip': master_ip, + 'work_queue_name': work_queue_name, + 'lithops_version': __version__, + } + + +def _mark_worker_error(worker_name, message): + """Records why a worker could not be set up, for the worker listing""" + redis_client.hset(_worker_key(worker_name), mapping={ + 'status': WorkerStatus.ERROR.value, + 'err': message, + }) + + +def _worker_setup_script(standalone_handler, vm_data): + """Returns the script that installs Lithops and the service on a worker""" + script = get_host_setup_script( + run_install=False, + **install_script_kwargs_from_config(standalone_handler.config), + ) + script += get_worker_setup_script(standalone_handler.config, vm_data) + return script + + def setup_worker_create_reuse(standalone_handler, worker_info, work_queue_name): """ - Run the worker setup process and installs all the Lithops dependencies into it + Installs Lithops on a worker VM, recreating the instance when it does not + come up or does not have what the runtime needs. The installation itself + is left running in the background, and the worker reports back when its + service comes up """ worker = standalone_handler.backend.get_instance(**worker_info, public=False) - if redis_client.hget(f"worker:{worker.name}", 'status') == WorkerStatus.ACTIVE.value: + if redis_client.hget(_worker_key(worker.name), 'status') == WorkerStatus.ACTIVE.value: return save_worker(worker, standalone_handler.config, work_queue_name) - max_instance_create_retries = worker.config.get('worker_create_retries', MAX_INSTANCE_CREATE_RETRIES) + max_retries = worker.config.get( + 'worker_create_retries', MAX_INSTANCE_CREATE_RETRIES + ) def wait_worker_ready(worker): + """Waits for a worker to boot, recreating it while there are tries left""" instance_ready_retries = 1 - - while instance_ready_retries <= max_instance_create_retries: + while instance_ready_retries <= max_retries: try: worker.wait_ready() break - except TimeoutError as e: # VM not started in time - redis_client.hset(f"worker:{worker.name}", 'status', WorkerStatus.ERROR.value) - err_msg = 'Timeout Error while waitting the VM to get ready' - redis_client.hset(f"worker:{worker.name}", 'err', err_msg) - if instance_ready_retries == max_instance_create_retries: + except TimeoutError: + err_msg = 'Timeout Error while waiting the VM to get ready' + _mark_worker_error(worker.name, err_msg) + if instance_ready_retries == max_retries: logger.debug(f'Readiness probe expired for {worker}') - raise e + raise logger.warning(f'Timeout Error. Recreating {worker}') worker.delete() worker.create() @@ -276,49 +369,40 @@ def wait_worker_ready(worker): wait_worker_ready(worker) instance_validate_retries = 1 - while instance_validate_retries <= max_instance_create_retries: + while instance_validate_retries <= max_retries: try: logger.debug(f'Validating {worker}') worker.validate_capabilities() break except LithopsValidationError as e: - redis_client.hset(f"worker:{worker.name}", 'status', WorkerStatus.ERROR.value) - redis_client.hset(f"worker:{worker.name}", 'err', f'Validation error: {e}') - if instance_validate_retries == max_instance_create_retries: + _mark_worker_error(worker.name, f'Validation error: {e}') + if instance_validate_retries == max_retries: logger.debug(f'Validation probe expired for {worker}') - raise e + raise logger.warning(f'{worker} validation error: {e}') worker.delete() worker.create() instance_validate_retries += 1 wait_worker_ready(worker) - redis_client.hset(f"worker:{worker.name}", 'private_ip', worker.private_ip) - redis_client.hset(f"worker:{worker.name}", 'status', WorkerStatus.STARTED.value) - redis_client.hset(f"worker:{worker.name}", 'err', '') + redis_client.hset(_worker_key(worker.name), mapping={ + 'private_ip': worker.private_ip, + 'status': WorkerStatus.STARTED.value, + 'err': '', + }) try: logger.debug(f'Uploading lithops files to {worker}') worker.get_ssh_client().upload_local_file( '/opt/lithops/lithops_standalone.zip', - '/tmp/lithops_standalone.zip') + '/tmp/lithops_standalone.zip', + ) logger.debug(f'Preparing installation script for {worker}') - vm_data = { - 'name': worker.name, - 'private_ip': worker.private_ip, - 'instance_id': worker.instance_id, - 'ssh_credentials': worker.ssh_credentials, - 'instance_type': worker.instance_type, - 'master_ip': master_ip, - 'work_queue_name': work_queue_name, - 'lithops_version': __version__ - } remote_script = "/tmp/install_lithops.sh" - script = get_host_setup_script( - run_install=False, **install_script_kwargs_from_config(standalone_handler.config) + script = _worker_setup_script( + standalone_handler, _worker_vm_data(worker, work_queue_name) ) - script += get_worker_setup_script(standalone_handler.config, vm_data) logger.debug(f'Submitting installation script to {worker}') worker.get_ssh_client().upload_data_to_file(script, remote_script) @@ -327,60 +411,64 @@ def wait_worker_ready(worker): worker.del_ssh_client() logger.debug(f'Installation script submitted to {worker}') - redis_client.hset(f"worker:{worker.name}", 'status', WorkerStatus.INSTALLING.value) - + redis_client.hset( + _worker_key(worker.name), 'status', WorkerStatus.INSTALLING.value + ) except Exception as e: - redis_client.hset(f"worker:{worker.name}", 'status', WorkerStatus.ERROR.value) - worker.err = f'Unable to setup lithops in the VM: {str(e)}' - raise e + _mark_worker_error( + worker.name, f'Unable to setup lithops in the VM: {e}' + ) + raise def setup_worker_consume(standalone_handler, worker_info, work_queue_name): """ - Run the worker setup process in the case of Consume mode + Installs the worker service on this very instance, which is what consume + mode runs the jobs on """ instance = standalone_handler.backend.get_instance(**worker_info, public=False) instance.private_ip = master_ip - if redis_client.hget(f"worker:{instance.name}", 'status') == WorkerStatus.ACTIVE.value: + if redis_client.hget(_worker_key(instance.name), 'status') == WorkerStatus.ACTIVE.value: return save_worker(instance, standalone_handler.config, work_queue_name) try: logger.debug(f'Setting up the worker in the current {instance}') - vm_data = { - 'name': instance.name, - 'private_ip': instance.private_ip, - 'instance_id': instance.instance_id, - 'ssh_credentials': instance.ssh_credentials, - 'instance_type': instance.instance_type, - 'master_ip': master_ip, - 'work_queue_name': work_queue_name, - 'lithops_version': __version__ - } worker_setup_script = "/tmp/install_lithops.sh" - script = get_host_setup_script( - run_install=False, **install_script_kwargs_from_config(standalone_handler.config) + script = _worker_setup_script( + standalone_handler, _worker_vm_data(instance, work_queue_name) ) - script += get_worker_setup_script(standalone_handler.config, vm_data) - with open(worker_setup_script, 'w') as wis: - wis.write(script) + with open(worker_setup_script, 'w') as script_file: + script_file.write(script) - redis_client.hset(f"worker:{instance.name}", 'status', WorkerStatus.INSTALLING.value) + redis_client.hset( + _worker_key(instance.name), + 'status', + WorkerStatus.INSTALLING.value, + ) os.chmod(worker_setup_script, 0o755) - os.system("sudo " + worker_setup_script) + # os.system reports the wait status, not the exit code, so it is + # logged as such rather than as a number that looks like one + wait_status = os.system("sudo " + worker_setup_script) + if wait_status != 0: + logger.error( + f'The setup script of {instance} failed with wait status ' + f'{wait_status}' + ) os.remove(worker_setup_script) - except Exception as e: - redis_client.hset(f"worker:{instance.name}", 'status', WorkerStatus.ERROR.value) - instance.err = f'Unable to setup lithops in the VM: {str(e)}' - raise e + _mark_worker_error( + instance.name, f'Unable to setup lithops in the VM: {e}' + ) + raise def handle_workers(job_payload, workers, work_queue_name): """ - Creates the workers (if any) + Sets up every worker of a job in parallel. A worker that fails to be set + up is one worker less, and the job runs on the ones that came up """ if not workers: return @@ -396,25 +484,20 @@ def handle_workers(job_payload, workers, work_queue_name): if standalone_config['exec_mode'] == StandaloneMode.CONSUME.value: try: setup_worker_consume( - standalone_handler, - workers[0], - work_queue_name + standalone_handler, workers[0], work_queue_name ) total_correct += 1 except Exception as e: - # TODO: If the local worker can't start, cancel all jobs - # in the budget keeper logger.error(e) else: with ThreadPoolExecutor(len(workers)) as executor: for worker_info in workers: - future = executor.submit( + futures.append(executor.submit( setup_worker_create_reuse, standalone_handler, worker_info, - work_queue_name - ) - futures.append(future) + work_queue_name, + )) for future in cf.as_completed(futures): try: @@ -435,16 +518,23 @@ def handle_workers(job_payload, workers, work_queue_name): def cancel_job_process(job_key_list): """ - Cleans the work queues and sends the SIGTERM to the workers + Cancels jobs: takes their tasks out of the work queue, tells every worker + to kill the ones already running, and marks the jobs as canceled """ for job_key in job_key_list: logger.debug(f'Received SIGTERM: Stopping job process {job_key}') - queue_name = redis_client.hget(f'job:{job_key}', 'queue_name') + queue_name = redis_client.hget(_job_key_id(job_key), 'queue_name') + if not queue_name: + logger.debug(f'Job {job_key} has no work queue to clean') + continue tmp_queue = [] while redis_client.llen(queue_name) > 0: task_payload_json = redis_client.rpop(queue_name) + if task_payload_json is None: + # A worker took the last task between the two calls + break task_payload = json.loads(task_payload_json) if task_payload['job_key'] != job_key: tmp_queue.append(task_payload_json) @@ -453,58 +543,68 @@ def cancel_job_process(job_key_list): redis_client.lpush(queue_name, task_payload_json) def stop_task(worker): + """Asks one worker to kill the tasks of this job""" worker_data = redis_client.hgetall(worker) - url = f"http://{worker_data['private_ip']}:{SA_WORKER_SERVICE_PORT}/stop/{job_key}" + url = ( + f"http://{worker_data['private_ip']}:" + f"{SA_WORKER_SERVICE_PORT}/stop/{job_key}" + ) requests.post(url, timeout=0.5) - # Send stop signal to all workers - workers = redis_client.keys('worker:*') - with ThreadPoolExecutor(len(workers)) as ex: - ex.map(stop_task, workers) + _map_if_any(stop_task, redis_client.keys('worker:*')) Path(os.path.join(JOBS_DIR, job_key + '.done')).touch() - if redis_client.hget(f"job:{job_key}", 'status') != JobStatus.DONE.value: - redis_client.hset(f"job:{job_key}", 'status', JobStatus.CANCELED.value) + if redis_client.hget(_job_key_id(job_key), 'status') != JobStatus.DONE.value: + redis_client.hset( + _job_key_id(job_key), 'status', JobStatus.CANCELED.value + ) @app.route('/job/stop', methods=['POST']) def stop(): - """ - Stops received job processes - """ - job_key_list = flask.request.get_json(force=True, silent=True) - # Start a separate thread to do the task in background, - # for not keeping the client waiting. - Thread(target=cancel_job_process, args=(job_key_list, )).start() - + """Cancels the given jobs, in the background""" + job_key_list = _json_body() + if not isinstance(job_key_list, list): + return error('The action did not receive a list as an argument.') + Thread(target=cancel_job_process, args=(job_key_list,)).start() return ('', 204) @app.route('/job/list', methods=['GET']) def list_jobs(): - """ - Returns the current workers state - """ + """Returns a table of every job the master knows about""" logger.debug('Listing jobs') - budget_keeper.last_usage_time = time.time() - result = [['Job ID', 'Function Name', 'Submitted', 'Worker Type', 'Runtime', 'Tasks Done', 'Job Status']] + result = [[ + 'Job ID', 'Function Name', 'Submitted', 'Worker Type', + 'Runtime', 'Tasks Done', 'Job Status', + ]] - for job_job_key in redis_client.keys('job:*'): - job_data = redis_client.hgetall(job_job_key) + for job_redis_key in redis_client.keys('job:*'): + job_data = redis_client.hgetall(job_redis_key) job_key = job_data['job_key'] exec_mode = job_data['exec_mode'] - status = job_data['status'] - func_name = job_data['func_name'] + "()" timestamp = float(job_data['submitted']) - runtime = job_data['runtime_name'] - worker_type = job_data['worker_type'] if exec_mode != StandaloneMode.CONSUME.value else 'VM' - submitted = datetime.fromtimestamp(timestamp).strftime('%Y-%m-%d %H:%M:%S UTC') + worker_type = ( + job_data['worker_type'] + if exec_mode != StandaloneMode.CONSUME.value + else 'VM' + ) + submitted = datetime.fromtimestamp(timestamp).strftime( + '%Y-%m-%d %H:%M:%S UTC' + ) total_tasks = str(job_data['total_tasks']) done_tasks = str(redis_client.llen(f'tasksdone:{job_key}')) - job = (job_key, func_name, submitted, worker_type, runtime, f'{done_tasks}/{total_tasks}', status) - result.append(job) + result.append(( + job_key, + job_data['func_name'] + "()", + submitted, + worker_type, + job_data['runtime_name'], + f'{done_tasks}/{total_tasks}', + job_data['status'], + )) logger.debug(f'jobs: {result}') return flask.jsonify(result) @@ -512,12 +612,12 @@ def list_jobs(): def handle_job(job_payload, queue_name): """ - Process responsible to put the job in redis and all the - individual tasks in a work queue + Registers a job and pushes one task per call into its work queue, each + task carrying only the data range of its own call """ job_key = job_payload['job_key'] - redis_client.hset(f"job:{job_key}", mapping={ + redis_client.hset(_job_key_id(job_key), mapping={ 'job_key': job_key, 'status': JobStatus.SUBMITTED.value, 'submitted': job_payload['host_submit_tstamp'], @@ -526,7 +626,7 @@ def handle_job(job_payload, queue_name): 'runtime_name': job_payload['runtime_name'], 'exec_mode': job_payload['config']['standalone']['exec_mode'], 'total_tasks': len(job_payload['call_ids']), - 'queue_name': queue_name + 'queue_name': queue_name, }) dbr = job_payload['data_byte_ranges'] @@ -536,17 +636,21 @@ def handle_job(job_payload, queue_name): task_payload['data_byte_ranges'] = [dbr[int(call_id)]] redis_client.lpush(queue_name, json.dumps(task_payload)) - logger.debug(f"Job {job_key} correctly submitted to work queue '{queue_name}'") + logger.debug( + f"Job {job_key} correctly submitted to work queue '{queue_name}'" + ) @app.route('/job/run', methods=['POST']) def run(): """ - Entry point for running jobs + Takes a job in: queues its tasks and sets its workers up, both in the + background, so that the caller is not left waiting for the VMs """ - job_payload = flask.request.get_json(force=True, silent=True) - if job_payload and not isinstance(job_payload, dict): - return error('The action did not receive a dictionary as an argument') + job_payload = _json_body() + bad_request = _require_dict(job_payload) + if bad_request is not None: + return bad_request try: runtime_name = job_payload['runtime_name'] @@ -559,8 +663,9 @@ def run(): budget_keeper.add_job(job_key) - exec_mode = job_payload['config']['standalone']['exec_mode'] - exec_mode = StandaloneMode[exec_mode.upper()] + exec_mode = StandaloneMode[ + job_payload['config']['standalone']['exec_mode'].upper() + ] workers = job_payload.pop('worker_instances') if exec_mode == StandaloneMode.CONSUME: @@ -570,7 +675,9 @@ def run(): elif exec_mode == StandaloneMode.REUSE: worker_it = job_payload['worker_instance_type'] worker_wp = job_payload['worker_processes'] - queue_name = f'wq:{worker_it}-{worker_wp}-{runtime_name.replace("/", "-")}'.lower() + queue_name = ( + f'wq:{worker_it}-{worker_wp}-{runtime_name.replace("/", "-")}' + ).lower() Thread(target=handle_job, args=(job_payload, queue_name)).start() Thread(target=handle_workers, args=(job_payload, workers, queue_name)).start() @@ -578,23 +685,28 @@ def run(): act_id = str(uuid.uuid4()).replace('-', '')[:12] response = flask.jsonify({'activationId': act_id}) response.status_code = 202 - return response def job_monitor(): + """ + Follows the tasks of every job as they finish, reporting the progress and + leaving behind the done file the budget keeper watches + """ logger.info("Starting job monitoring thread") - jobs_data = {} while True: time.sleep(JOB_MONITOR_CHECK_INTERVAL) - for job_job_key in redis_client.keys('job:*'): - job_key = job_job_key.replace("job:", "") + for job_redis_key in redis_client.keys('job:*'): + job_key = job_redis_key.replace("job:", "") if job_key not in jobs_data: budget_keeper.add_job(job_key) - job_data = redis_client.hgetall(job_job_key) - jobs_data[job_key] = {'total': int(job_data['total_tasks']), 'done': 0} + job_data = redis_client.hgetall(job_redis_key) + jobs_data[job_key] = { + 'total': int(job_data['total_tasks']), + 'done': 0, + } if jobs_data[job_key]['total'] == jobs_data[job_key]['done']: continue done_tasks = int(redis_client.llen(f"tasksdone:{job_key}")) @@ -602,7 +714,10 @@ def job_monitor(): total_tasks = jobs_data[job_key]['total'] jobs_data[job_key]['done'] = done_tasks exec_id, job_id = job_key.rsplit('-', 1) - msg = f"ExecutorID: {exec_id} | JObID: {job_id} - Tasks done: {done_tasks}/{total_tasks}" + msg = ( + f'{log_prefix(exec_id, job_id)} - ' + f'Tasks done: {done_tasks}/{total_tasks}' + ) if jobs_data[job_key]['total'] == jobs_data[job_key]['done']: Path(os.path.join(JOBS_DIR, f'{job_key}.done')).touch() msg += " - Completed!" @@ -615,20 +730,22 @@ def job_monitor(): @app.route('/clean', methods=['POST']) def clean(): + """Drops every job and worker the master had recorded""" logger.debug("Clean command received. Cleaning all data from redis") redis_client.flushall() - return ('', 204) @app.route('/ping', methods=['GET']) def ping(): - response = flask.jsonify({'response': lithops_version}) + """Answers with the Lithops version this master runs""" + response = flask.jsonify({'response': __version__}) response.status_code = 200 return response def error(msg): + """Builds the response of a request the master could not act on""" response = flask.jsonify({'error': msg}) response.status_code = 404 return response @@ -636,50 +753,64 @@ def error(msg): @app.route('/metadata', methods=['GET']) def get_metadata(): - payload = flask.request.get_json(force=True, silent=True) - if payload and not isinstance(payload, dict): - return error('The action did not receive a dictionary as an argument.') + """ + Returns the metadata of a runtime, which the master extracts by running + it locally, as it is the only instance that is up at this point + """ + payload = _json_body() + bad_request = _require_dict(payload) + if bad_request is not None: + return bad_request try: verify_runtime_name(payload['runtime']) except Exception as e: return error(str(e)) - localhos_handler = LocalhostHandler(payload) - localhos_handler.init() - runtime_meta = localhos_handler.deploy_runtime(payload['runtime']) + localhost_handler = LocalhostHandler(payload) + localhost_handler.init() + runtime_meta = localhost_handler.deploy_runtime(payload['runtime']) if 'lithops_version' in runtime_meta: - logger.debug(f"Runtime metdata extracted correctly from {payload['runtime']}" - f" - Lithops {runtime_meta['lithops_version']}") + logger.debug( + f"Runtime metadata extracted correctly from {payload['runtime']}" + f" - Lithops {runtime_meta['lithops_version']}" + ) response = flask.jsonify(runtime_meta) response.status_code = 200 - return response def main(): + """ + Entry point of the master service: starts the countdown that stops the + instance, the job monitor, and the endpoints the client talks to + """ global redis_client global budget_keeper global master_ip - os.makedirs(LITHOPS_TEMP_DIR, exist_ok=True) + _configure_logging() - with open(SA_CONFIG_FILE, 'r') as cf: - standalone_config = json.load(cf) + with open(SA_CONFIG_FILE, 'r') as config_file: + standalone_config = json.load(config_file) - with open(SA_MASTER_DATA_FILE, 'r') as ad: - master_data = json.load(ad) + with open(SA_MASTER_DATA_FILE, 'r') as data_file: + master_data = json.load(data_file) master_ip = master_data['private_ip'] - budget_keeper = BudgetKeeper(standalone_config, master_data, stop_callback=clean) + budget_keeper = BudgetKeeper( + standalone_config, master_data, stop_callback=clean + ) budget_keeper.start() redis_client = redis.Redis(decode_responses=True) Thread(target=job_monitor, daemon=True).start() - server = WSGIServer(('0.0.0.0', SA_MASTER_SERVICE_PORT), app, log=app.logger) + server = WSGIServer( + ('0.0.0.0', SA_MASTER_SERVICE_PORT), app, log=app.logger + ) server.serve_forever() diff --git a/lithops/standalone/runner.py b/lithops/standalone/runner.py index 7ff764d96..a620fe6ee 100644 --- a/lithops/standalone/runner.py +++ b/lithops/standalone/runner.py @@ -21,27 +21,44 @@ import uuid from lithops.worker import function_handler -from lithops.constants import ( - RN_LOG_FILE, - LOGGER_FORMAT -) +from lithops.utils import log_prefix +from lithops.constants import RN_LOG_FILE, LITHOPS_TEMP_DIR, LOGGER_FORMAT -log_file_stream = open(RN_LOG_FILE, 'a') -logging.basicConfig(stream=log_file_stream, level=logging.INFO, format=LOGGER_FORMAT) logger = logging.getLogger('lithops.standalone.runner') -def run_job(backend, task_filename): +def _configure_runner_logging(): + """ + Sends everything this process logs to the runner log, which is the only + place a task failure can be read from in a standalone worker + """ + os.makedirs(LITHOPS_TEMP_DIR, exist_ok=True) + log_file_stream = open(RN_LOG_FILE, 'a') + logging.basicConfig( + stream=log_file_stream, + level=logging.INFO, + format=LOGGER_FORMAT, + ) + return log_file_stream + + +def run_job(backend: str, task_filename: str) -> None: + """ + Runs the single task described by a task file. The worker service starts + one of these per task, so the handler always runs one activation + """ logger.info(f'Got {task_filename} job file') - with open(task_filename, 'rb') as jf: - task_payload = json.load(jf) + with open(task_filename, 'r') as task_file: + task_payload = json.load(task_file) executor_id = task_payload['executor_id'] job_id = task_payload['job_id'] call_id = task_payload['call_ids'][0] - logger.info(f'ExecutorID {executor_id} | JobID {job_id} | CallID {call_id} - Starting execution') + logger.info( + f'{log_prefix(executor_id, job_id, call_id)} - Starting execution' + ) act_id = str(uuid.uuid4()).replace('-', '')[:12] os.environ['__LITHOPS_ACTIVATION_ID'] = act_id @@ -50,14 +67,24 @@ def run_job(backend, task_filename): task_payload['worker_processes'] = 1 function_handler(task_payload) - logger.info(f'ExecutorID {executor_id} | JobID {job_id} | CallID {call_id} - Execution Finished') + logger.info( + f'{log_prefix(executor_id, job_id, call_id)} - Execution Finished' + ) + + +def main() -> None: + """Entry point of the task runner, called by the worker service""" + log_file_stream = _configure_runner_logging() + try: + sys.stdout = log_file_stream + sys.stderr = log_file_stream + logger.info('Starting Standalone task runner') + backend = sys.argv[1] + task_filename = sys.argv[2] + run_job(backend, task_filename) + finally: + log_file_stream.close() if __name__ == "__main__": - sys.stdout = log_file_stream - sys.stderr = log_file_stream - logger.info('Starting Standalone task runner') - backend = sys.argv[1] - task_filename = sys.argv[2] - run_job(backend, task_filename) - log_file_stream.close() + main() diff --git a/lithops/standalone/standalone.py b/lithops/standalone/standalone.py index b969ab067..7831c39bb 100644 --- a/lithops/standalone/standalone.py +++ b/lithops/standalone/standalone.py @@ -25,11 +25,13 @@ import requests import shlex import concurrent.futures as cf +from typing import Any, Dict from lithops.utils import ( BackendType, is_lithops_worker, - create_handler_zip + create_handler_zip, + log_prefix, ) from lithops.constants import ( TEMP_DIR, @@ -47,14 +49,17 @@ logger = logging.getLogger(__name__) +_CURL_BODY_INLINE_LIMIT = 130000 + class StandaloneHandler: """ - A StandaloneHandler object is used by invokers and other components to access - underlying standalone backend without exposing the implementation details. + A StandaloneHandler object is used by invokers and other components to + access the underlying standalone backend without exposing implementation + details. """ - def __init__(self, standalone_config): + def __init__(self, standalone_config: Dict[str, Any]): self.config = standalone_config self.backend_name = self.config['backend'] self.start_timeout = self.config['start_timeout'] @@ -63,261 +68,369 @@ def __init__(self, standalone_config): module_location = f'lithops.standalone.backends.{self.backend_name}' sb_module = importlib.import_module(module_location) - StandaloneBackend = getattr(sb_module, 'StandaloneBackend') - self.backend = StandaloneBackend(self.config[self.backend_name], self.exec_mode.value) + standalone_backend_cls = getattr(sb_module, 'StandaloneBackend') + self.backend = standalone_backend_cls( + self.config[self.backend_name], self.exec_mode.value + ) - self.jobs = [] # list to store executed jobs (job_keys) + self.jobs = [] logger.debug("Standalone handler created successfully") def init(self): - """ - Initialize the backend and create/start the master VM instance - """ + """Prepares the backend resources a standalone run needs""" self.backend.init() def is_initialized(self): - """ - Check if the backend is initialized - """ + """True when the backend has resources from a previous run""" return self.backend.is_initialized() - def build_image(self, image_name, script_file, overwrite, include, extra_args=[]): - """ - Builds a new VM Image - """ - self.backend.build_image(image_name, script_file, overwrite, include, extra_args) + def build_image( + self, image_name, script_file, overwrite, include, extra_args=None + ): + """Builds the VM image the instances will boot from""" + self.backend.build_image( + image_name, script_file, overwrite, include, extra_args or [] + ) def delete_image(self, name): - """ - Deletes VM Image - """ + """Deletes a VM image built for Lithops""" self.backend.delete_image(name) def list_images(self): + """Lists the VM images built for Lithops""" + return self.backend.list_images() + + def _master_url(self, endpoint: str) -> str: """ - Lists VM Images + Returns the URL of a master service endpoint. A worker reaches the + master by the name the setup script put in its hosts file """ - return self.backend.list_images() + host = 'lithops-master' if self.is_lithops_worker else '127.0.0.1' + return f'http://{host}:{SA_MASTER_SERVICE_PORT}/{endpoint}' - def _make_request(self, method, endpoint, data=None): + def _make_request(self, method: str, endpoint: str, data=None): """ - Makes a requests to the master VM + Calls the master service. A worker can reach it over the network, + while the client has to go through the SSH connection it already has """ if self.is_lithops_worker: - url = f"http://lithops-master:{SA_MASTER_SERVICE_PORT}/{endpoint}" - if method == 'GET': - resp = requests.get(url, timeout=1) - return resp.json() - elif method == 'POST': - resp = requests.post(url, data=json.dumps(data)) - resp.raise_for_status() - return resp.json() - else: - url = f'http://127.0.0.1:{SA_MASTER_SERVICE_PORT}/{endpoint}' - cmd = f'curl -X {method} {url} -H \'Content-Type: application/json\'' - if data is not None: - json_data = json.dumps(data) - data_size = len(json_data) - if data_size < 130000: - data_str = shlex.quote(json_data) - cmd = f'{cmd} -d {data_str}' - else: - data_file_name = f'/tmp/lithops_data_{str(uuid.uuid4())[-6:]}.json' - self.backend.master.get_ssh_client().upload_data_to_file(json_data, data_file_name) - cmd = f'{cmd} -d @{data_file_name}; rm {data_file_name}' - out, err = self.backend.master.get_ssh_client().run_remote_command(cmd) - if not out: + return self._request_from_worker(method, endpoint, data) + return self._request_via_ssh(method, endpoint, data) + + def _request_from_worker(self, method: str, endpoint: str, data=None): + """Calls the master service over HTTP, from inside the network""" + url = self._master_url(endpoint) + if method == 'GET': + resp = requests.get(url, timeout=1) + return resp.json() + if method == 'POST': + resp = requests.post(url, data=json.dumps(data)) + resp.raise_for_status() + if not resp.content: + return None + return resp.json() + raise ValueError(f'Unsupported HTTP method: {method}') + + def _request_via_ssh(self, method: str, endpoint: str, data=None): + """ + Calls the master service through curl over SSH. A body too large for a + command line is uploaded to a file the remote curl then reads + """ + url = self._master_url(endpoint) + # -sS hides the progress meter. Without it, a 204 from /job/stop or + # /clean leaves stdout empty and the meter on stderr, which used to + # be raised as "Could not stop the jobs on the master" + cmd = ( + f"curl -sS -X {method} {url} " + f"-H 'Content-Type: application/json'" + ) + if data is not None: + json_data = json.dumps(data) + if len(json_data) < _CURL_BODY_INLINE_LIMIT: + cmd = f'{cmd} -d {shlex.quote(json_data)}' + else: + data_file_name = ( + f'/tmp/lithops_data_{str(uuid.uuid4())[-6:]}.json' + ) + self.backend.master.get_ssh_client().upload_data_to_file( + json_data, data_file_name + ) + cmd = f'{cmd} -d @{data_file_name}; rm {data_file_name}' + + out, err = self.backend.master.get_ssh_client().run_remote_command(cmd) + if not out: + if err: raise ValueError(err) - try: - return json.loads(out) - except Exception: - raise ValueError(out) + return None + try: + return json.loads(out) + except Exception as e: + # Whatever the master printed instead of a response is the only + # clue about what went wrong there + raise ValueError(out) from e def _is_master_service_ready(self): """ - Checks if the proxy is ready to receive http connections + True when the master service answers and runs this same Lithops + version, as a master left over from another version cannot be trusted """ try: resp = self._make_request('GET', 'ping') if resp['response'] != __version__: raise LithopsValidationError( - f"{self.backend.master} is running Lithops {resp['response']} and " - f"it doesn't match local lithops version {__version__}, consider running " - f"'lithops clean -b {self.backend_name} --all' to delete the master instance") + f"{self.backend.master} is running Lithops " + f"{resp['response']} and it doesn't match local lithops " + f"version {__version__}, consider running " + f"'lithops clean -b {self.backend_name} --all' to delete " + f"the master instance" + ) return True - except LithopsValidationError as e: - raise e + except LithopsValidationError: + raise except Exception: return False def _validate_master_service_setup(self): """ - Checks the master VM is correctly installed + Makes sure the master has the service installed and running, setting + it up when it was never installed and giving up when it is dead """ - logger.debug(f'Validating lithops master service is installed on {self.backend.master}') + logger.debug( + f'Validating lithops master service is installed on ' + f'{self.backend.master}' + ) ssh_client = self.backend.master.get_ssh_client() out, err = ssh_client.run_remote_command(f'cat {SA_MASTER_DATA_FILE}') if not out: self._setup_master_service() return - logger.debug(f"Validating lithops master service is running on {self.backend.master}") + logger.debug( + f"Validating lithops master service is running on " + f"{self.backend.master}" + ) out, err = ssh_client.run_remote_command("service lithops-master status") if not out or 'Active: active (running)' not in out: self.dismantle() raise LithopsValidationError( f"Lithops master service not active on {self.backend.master}, " "consider to delete master instance and metadata using " - "'lithops clean --all'") + "'lithops clean --all'" + ) def _wait_master_service_ready(self): - """ - Waits until the master service is ready to receive http connections - """ - logger.info(f'Waiting for Lithops service to become ready on {self.backend.master}') + """Waits until the master service answers, or gives the instance up""" + logger.info( + f'Waiting for Lithops service to become ready on ' + f'{self.backend.master}' + ) start = time.time() - while (time.time() - start < self.start_timeout): + while time.time() - start < self.start_timeout: if self._is_master_service_ready(): ready_time = round(time.time() - start, 2) - logger.debug(f'{self.backend.master} ready in {ready_time} seconds') + logger.debug( + f'{self.backend.master} ready in {ready_time} seconds' + ) return True time.sleep(2) self.dismantle() - raise Exception(f'Lithops service readiness probe expired on {self.backend.master}') + raise Exception( + f'Lithops service readiness probe expired on {self.backend.master}' + ) - def _get_workers_on_master(self, worker_instance_type, worker_processes, runtime_name): + def _get_workers_on_master( + self, worker_instance_type, worker_processes, runtime_name + ): """ - gets the total available workers on the master VM + Returns the free workers the master already has of the requested + shape, and none when it cannot be asked """ - workers_on_master = [] try: payload = { 'worker_instance_type': worker_instance_type, 'worker_processes': worker_processes, - 'runtime_name': runtime_name + 'runtime_name': runtime_name, } - workers_on_master = self._make_request('GET', 'worker/get', payload) - except Exception: - pass - return workers_on_master + return self._make_request('GET', 'worker/get', payload) + except Exception as e: + logger.debug(f'Could not get the workers of the master: {e}') + return [] + + def _create_workers(self, workers_to_create: int, executor_id, job_id): + """ + Creates worker instances in parallel and returns the ones that came + up. A worker that fails to be created is one worker less, not a + failed job, so the job runs on whatever came up + """ + if workers_to_create <= 0: + return [] + current_workers_old = set(self.backend.workers) + futures = [] + with cf.ThreadPoolExecutor(min(workers_to_create, 48)) as ex: + for vm_n in range(workers_to_create): + worker_id = f"{executor_id}-{job_id}-{vm_n}" + worker_hash = hashlib.sha1( + worker_id.encode("utf-8") + ).hexdigest()[:8] + name = f'lithops-worker-{worker_hash}' + futures.append(ex.submit(self.backend.create_worker, name)) + + for future in cf.as_completed(futures): + try: + future.result() + except Exception as e: + logger.debug(f'Could not create a worker instance: {e}') - def invoke(self, job_payload): + new_workers = set(self.backend.workers) - current_workers_old + logger.debug( + f"Total worker VM instances created: " + f"{len(new_workers)}/{workers_to_create}" + ) + return list(new_workers) + + def _required_workers(self, job_payload) -> int: """ - Run the job description against the selected environment + Returns how many workers the job needs, filling in the instance shape + the backend offers. Capped by max_workers, as that is the limit the + user set on how much the run may spend """ executor_id = job_payload['executor_id'] job_id = job_payload['job_id'] total_calls = job_payload['total_calls'] - if self.exec_mode == StandaloneMode.CONSUME: - logger.debug( - f'ExecutorID {executor_id} | JobID {job_id} - Worker processes: ' - f'{job_payload["worker_processes"]}' - ) - else: - worker_instance_type = self.backend.get_worker_instance_type() - worker_processes = self.backend.get_worker_cpu_count() + worker_instance_type = self.backend.get_worker_instance_type() + worker_processes = self.backend.get_worker_cpu_count() + job_payload['worker_instance_type'] = worker_instance_type - job_payload['worker_instance_type'] = worker_instance_type + if job_payload['worker_processes'] == "AUTO": + job_payload['worker_processes'] = worker_processes + job_payload['config'][self.backend_name]['worker_processes'] = ( + worker_processes + ) - if job_payload['worker_processes'] == "AUTO": - job_payload['worker_processes'] = worker_processes - job_payload['config'][self.backend_name]['worker_processes'] = worker_processes + wp = job_payload['worker_processes'] + max_workers = job_payload['max_workers'] + required_workers = min( + max_workers, total_calls // wp + (total_calls % wp > 0) + ) + logger.debug( + f'{log_prefix(executor_id, job_id)} - Instance Type: ' + f'{worker_instance_type} - Worker ' + f'processes: {job_payload["worker_processes"]} - ' + f'Required Workers: {required_workers}' + ) + return required_workers + + def _acquire_workers(self, job_payload, required_workers: int): + """ + Returns the workers the job will run on, as the instances that were + created for it and the total the job can count on. Consume mode runs + on the master itself, and reuse mode only creates what the master does + not already have free + """ + executor_id = job_payload['executor_id'] + job_id = job_payload['job_id'] - wp = job_payload['worker_processes'] - max_workers = job_payload['max_workers'] - required_workers = min(max_workers, total_calls // wp + (total_calls % wp > 0)) + if self.exec_mode == StandaloneMode.CONSUME: + return [self.backend.master], 1 - logger.debug( - f'ExecutorID {executor_id} | JobID {job_id} - Instance Type: {worker_instance_type} - Worker ' - f'processes: {job_payload["worker_processes"]} - Required Workers: {required_workers}' + if self.exec_mode == StandaloneMode.CREATE: + new_workers = self._create_workers( + required_workers, executor_id, job_id ) + return new_workers, len(new_workers) + + workers = self._get_workers_on_master( + job_payload['worker_instance_type'], + job_payload['worker_processes'], + job_payload['runtime_name'], + ) + total_workers = len(workers) + logger.debug( + f"Found {total_workers} free workers connected to " + f"{self.backend.master}" + ) - def create_workers(workers_to_create): - current_workers_old = set(self.backend.workers) - futures = [] - with cf.ThreadPoolExecutor(min(workers_to_create, 48)) as ex: - for vm_n in range(workers_to_create): - worker_id = f"{executor_id}-{job_id}-{vm_n}" - worker_hash = hashlib.sha1(worker_id.encode("utf-8")).hexdigest()[:8] - name = f'lithops-worker-{worker_hash}' - futures.append(ex.submit(self.backend.create_worker, name)) + new_workers = [] + if total_workers < required_workers: + workers_to_create = required_workers - total_workers + logger.debug(f'Going to create {workers_to_create} new workers') + new_workers = self._create_workers( + workers_to_create, executor_id, job_id + ) + total_workers += len(new_workers) - for future in cf.as_completed(futures): - try: - future.result() - except Exception: - pass + return new_workers, total_workers - current_workers_new = set(self.backend.workers) - new_workers = current_workers_new - current_workers_old - logger.debug(f"Total worker VM instances created: {len(new_workers)}/{workers_to_create}") + def _ensure_master_ready(self) -> None: + """Sets the master service up unless it is already answering""" + logger.debug(f"Checking if {self.backend.master} is ready") + if self._is_master_service_ready(): + return - return list(new_workers) + self.backend.master.create(check_if_exists=True) + self.backend.master.wait_ready() + self._validate_master_service_setup() + self._wait_master_service_ready() - new_workers = [] + def invoke(self, job_payload): + """ + Runs a job on the standalone backend: works out how many workers it + needs, gets them up, and hands the job over to the master service + """ + executor_id = job_payload['executor_id'] + job_id = job_payload['job_id'] + total_calls = job_payload['total_calls'] + required_workers = 0 if self.exec_mode == StandaloneMode.CONSUME: - new_workers.append(self.backend.master) - total_workers = 1 - - elif self.exec_mode == StandaloneMode.CREATE: - new_workers = create_workers(required_workers) - total_workers = len(new_workers) - - elif self.exec_mode == StandaloneMode.REUSE: - workers = self._get_workers_on_master( - job_payload['worker_instance_type'], - job_payload['worker_processes'], - job_payload['runtime_name'], + logger.debug( + f'{log_prefix(executor_id, job_id)} - Worker processes: ' + f'{job_payload["worker_processes"]}' ) - total_workers = len(workers) - logger.debug(f"Found {total_workers} free workers connected to {self.backend.master}") - if total_workers < required_workers: - # create missing delta of workers - workers_to_create = required_workers - total_workers - logger.debug(f'Going to create {workers_to_create} new workers') - new_workers = create_workers(workers_to_create) - total_workers += len(new_workers) + else: + required_workers = self._required_workers(job_payload) + + new_workers, total_workers = self._acquire_workers( + job_payload, required_workers + ) if total_workers == 0: raise Exception('It was not possible to create any workers') - logger.debug(f'ExecutorID {executor_id} | JobID {job_id} - Going to run ' - f'{total_calls} activations in {total_workers} workers') + logger.debug( + f'{log_prefix(executor_id, job_id)} - Going to run ' + f'{total_calls} activations in {total_workers} workers' + ) - logger.debug(f"Checking if {self.backend.master} is ready") - if not self._is_master_service_ready(): - self.backend.master.create(check_if_exists=True) - self.backend.master.wait_ready() - self._validate_master_service_setup() - self._wait_master_service_ready() + self._ensure_master_ready() - # delete ssh key + # The key never leaves the client: the master has its own to reach the + # workers it creates backend = job_payload['config']['lithops']['backend'] job_payload['config'][backend].pop('ssh_key_filename', None) - # prepare worker instances data job_payload['worker_instances'] = [ - {'name': inst.name, - 'private_ip': inst.private_ip, - 'instance_id': inst.instance_id, - 'ssh_credentials': inst.ssh_credentials, - 'instance_type': inst.instance_type} + { + 'name': inst.name, + 'private_ip': inst.private_ip, + 'instance_id': inst.instance_id, + 'ssh_credentials': inst.ssh_credentials, + 'instance_type': inst.instance_type, + } for inst in new_workers ] - # invoke Job self._make_request('POST', 'job/run', job_payload) logger.debug(f'Job invoked on {self.backend.master}') - self.jobs.append(job_payload['job_key']) def deploy_runtime(self, runtime_name, *args): """ - Installs the proxy and extracts the runtime metadata + Brings the master up, installs the service on it, and asks it for the + metadata of the runtime, which only the master can extract """ logger.debug(f'Checking if {self.backend.master} is ready') if not self.backend.master.is_ready(): @@ -329,97 +442,96 @@ def deploy_runtime(self, runtime_name, *args): logger.debug('Extracting runtime metadata information') payload = {'runtime': runtime_name, 'pull_runtime': True} - runtime_meta = self._make_request('GET', 'metadata', payload) - - return runtime_meta + return self._make_request('GET', 'metadata', payload) def dismantle(self, **kwargs): - """ - Stop all VM instances - """ + """Stops the instances of this run""" self.backend.dismantle(**kwargs) def clean(self, **kwargs): """ - Clan all the backend resources + Deletes the resources of this run. The master is asked to clean up + after itself first, unless everything is going away anyway """ all_clean = kwargs.get('all', False) if self.is_initialized() and not all_clean: try: self.init() self._make_request('POST', 'clean') - except Exception: - pass + except Exception as e: + # A master that cannot be reached has nothing to clean, and + # the backend cleanup below still runs + logger.debug(f'Could not clean up through the master: {e}') self.backend.clean(**kwargs) def clear(self, job_keys=None, exception=None): """ - Clear all the backend resources. - clear method is executed after the results are get, - when an exception is produced, or when a user press ctrl+c + Stops the jobs this handler invoked. Workers meant to be reused stay + up, as the next job is going to run on them """ try: self._make_request('POST', 'job/stop', self.jobs) - except Exception: - pass + logger.debug('Jobs stopped on the master') + except Exception as e: + logger.debug(f'Could not stop the jobs on the master: {e}') if self.exec_mode != StandaloneMode.REUSE: self.backend.clear(job_keys) def list_jobs(self): - """ - Lists jobs in master VM - """ + """Lists the jobs the master knows about""" return self._make_request('GET', 'job/list') def list_workers(self): - """ - Lists available workers in master VM - """ + """Lists the workers connected to the master""" return self._make_request('GET', 'worker/list') - def get_runtime_key(self, runtime_name, runtime_memory, version=__version__): + def get_runtime_key( + self, runtime_name, runtime_memory, version=__version__ + ): """ - Wrapper method that returns a formated string that represents the - runtime key. Each backend has its own runtime key format. Used to - store runtime metadata into the storage + Returns a formatted string that represents the runtime key. + Each backend has its own runtime key format. Used to store + runtime metadata in storage. """ return self.backend.get_runtime_key(runtime_name, version) def get_runtime_info(self): - """ - Method that returns a dictionary with all the runtime information - set in config - """ - runtime_info = { + """Returns the runtime limits the executor reports to the user""" + return { 'runtime_name': self.config['runtime'], 'runtime_memory': None, 'runtime_timeout': self.config['hard_dismantle_timeout'], 'max_workers': self.config[self.backend_name]['max_workers'], } - return runtime_info - def get_backend_type(self): - """ - Wrapper method that returns the type of the backend (Batch or FaaS) - """ + """Returns the backend type, which is invoked with a whole job""" return BackendType.BATCH.value def _setup_master_service(self): """ - Setup lithops necessary packages and files in master VM instance + Installs Lithops on the master and starts its service, then brings + back the public key the master generated, which is what the workers + it creates will trust """ logger.info(f'Installing Lithops in {self.backend.master}') ssh_client = self.backend.master.get_ssh_client() - handler_zip = os.path.join(TEMP_DIR, f'lithops_standalone_{str(uuid.uuid4())[-6:]}.zip') - worker_path = os.path.join(os.path.dirname(__file__), 'worker.py') - master_path = os.path.join(os.path.dirname(__file__), 'master.py') - runner_path = os.path.join(os.path.dirname(__file__), 'runner.py') - create_handler_zip(handler_zip, [master_path, worker_path, runner_path]) + handler_zip = os.path.join( + TEMP_DIR, f'lithops_standalone_{str(uuid.uuid4())[-6:]}.zip' + ) + module_dir = os.path.dirname(__file__) + create_handler_zip( + handler_zip, + [ + os.path.join(module_dir, 'master.py'), + os.path.join(module_dir, 'worker.py'), + os.path.join(module_dir, 'runner.py'), + ], + ) logger.debug(f'Uploading lithops files to {self.backend.master}') ssh_client.upload_local_file(handler_zip, '/tmp/lithops_standalone.zip') @@ -430,14 +542,21 @@ def _setup_master_service(self): 'instance_id': self.backend.master.get_instance_id(), 'private_ip': self.backend.master.get_private_ip(), 'delete_on_dismantle': self.backend.master.delete_on_dismantle, - 'lithops_version': __version__ + 'lithops_version': __version__, } - logger.debug(f'Executing lithops installation process on {self.backend.master}') - logger.debug('Be patient, initial installation process may take up to 3 minutes') + logger.debug( + f'Executing lithops installation process on {self.backend.master}' + ) + logger.debug( + 'Be patient, initial installation process may take up to 3 minutes' + ) remote_script = "/tmp/install_lithops.sh" - script = get_host_setup_script(run_install=False, **install_script_kwargs_from_config(self.config)) + script = get_host_setup_script( + run_install=False, + **install_script_kwargs_from_config(self.config), + ) script += get_master_setup_script(self.config, master_data) ssh_client.upload_data_to_file(script, remote_script) @@ -448,4 +567,5 @@ def _setup_master_service(self): # This public key will be used to create the workers ssh_client.download_remote_file( f'{self.backend.master.home_dir}/.ssh/lithops_id_rsa.pub', - f'{self.backend.cache_dir}/{self.backend.master.name}-id_rsa.pub') + f'{self.backend.cache_dir}/{self.backend.master.name}-id_rsa.pub', + ) diff --git a/lithops/standalone/utils.py b/lithops/standalone/utils.py index aa481fc5e..13bd38799 100644 --- a/lithops/standalone/utils.py +++ b/lithops/standalone/utils.py @@ -3,7 +3,9 @@ import json import shlex from enum import Enum +from typing import Any, Dict, List, Tuple +from lithops.localhost.config import LocalhostEnvironment, get_environment from lithops.constants import ( SA_INSTALL_DIR, SA_SETUP_LOG_FILE, @@ -17,27 +19,32 @@ class StandaloneMode(Enum): + """ + How a standalone run uses its VMs: run everything on the master, create + one set of workers per job, or keep the workers around for the next job + """ + CONSUME = "consume" CREATE = "create" REUSE = "reuse" -def prepare_standalone_clean(backend, load_cache_fn): +def prepare_standalone_clean(backend, load_cache_fn) -> None: """ - Load persisted stack metadata from disk when the backend has a cache file. - - Standalone cloud backends call this at the start of clean() so cleanup works - even when clean() is invoked without a prior init() in the same process. + Loads the stack metadata a previous run persisted on disk, so that clean() + works even when it is called without an init() in the same process """ if backend.is_initialized(): load_cache_fn() -def standalone_clean_stop_early(backend, stack_data, delete_cache_fn, all_flag): +def standalone_clean_stop_early( + backend, stack_data, delete_cache_fn, all_flag +) -> bool: """ - Common clean() early exits for consume mode and missing stack metadata. - - Returns True when no further cloud resource cleanup is required. + Handles the clean() cases that own no cloud resources: consume mode, which + runs on an instance the user manages, and a stack nothing was created for. + Returns True when there is nothing else to clean """ if backend.mode == StandaloneMode.CONSUME.value: delete_cache_fn() @@ -50,6 +57,8 @@ def standalone_clean_stop_early(backend, stack_data, delete_cache_fn, all_flag): class WorkerStatus(Enum): + """States a worker VM reports while it is being set up and used""" + STARTING = "starting" STARTED = "started" ERROR = "error" @@ -61,6 +70,8 @@ class WorkerStatus(Enum): class JobStatus(Enum): + """States a job goes through in a standalone run""" + SUBMITTED = "submitted" PENDING = "pending" RUNNING = "running" @@ -69,7 +80,12 @@ class JobStatus(Enum): class LithopsValidationError(Exception): - pass + """Raised when the setup of a standalone run cannot be trusted""" + + +def is_container_runtime(runtime_name: str) -> bool: + """True when the runtime is a container image and not an interpreter""" + return get_environment(runtime_name) is LocalhostEnvironment.CONTAINER MASTER_SERVICE_NAME = 'lithops-master.service' @@ -116,7 +132,7 @@ class LithopsValidationError(Exception): shell: /bin/bash """ -CLOUD_CONFIG_WORKER = """ +CLOUD_CONFIG_WORKER = r""" #cloud-config bootcmd: - echo '{0}:{1}' | chpasswd @@ -130,7 +146,11 @@ class LithopsValidationError(Exception): """ -def _normalize_package_list(packages): +def _normalize_package_list(packages) -> List[str]: + """ + Returns the packages of a config entry as a list, accepting both a list + and a space separated string + """ if not packages: return [] if isinstance(packages, str): @@ -138,7 +158,12 @@ def _normalize_package_list(packages): return [str(p).strip() for p in packages if str(p).strip()] -def _format_apt_packages_for_shell(packages): +def _format_apt_packages_for_shell(packages) -> str: + """ + Returns the apt packages as one argument list for the setup script. The + names go into a shell command, so anything that is not a package name is + rejected instead of quoted + """ safe = [] for package in _normalize_package_list(packages): if not re.match(r'^[a-z0-9][a-z0-9.+~-]*$', package, re.IGNORECASE): @@ -149,7 +174,12 @@ def _format_apt_packages_for_shell(packages): return ' '.join(safe) -def _format_pip_packages_for_shell(packages): +def _format_pip_packages_for_shell(packages) -> str: + """ + Returns the pip specs as one argument list for the setup script. Specs + carry version markers, so they are quoted rather than restricted, and only + shell metacharacters are rejected + """ quoted = [] for package in _normalize_package_list(packages): if re.search(r'[;&|`$(){}]', package): @@ -160,22 +190,28 @@ def _format_pip_packages_for_shell(packages): return ' '.join(quoted) -def install_script_kwargs_from_config(config=None): +def install_script_kwargs_from_config(config=None) -> Dict[str, str]: """ - Build keyword arguments for get_host_setup_script() from standalone config. + Returns the arguments get_host_setup_script() takes, read from the + standalone configuration """ config = config or {} return { 'lithops_pip_spec': lithops_pip_spec_from_config(config), - 'extra_apt_packages': _format_apt_packages_for_shell(config.get('extra_apt_packages')), - 'extra_python_packages': _format_pip_packages_for_shell(config.get('extra_python_packages')), + 'extra_apt_packages': _format_apt_packages_for_shell( + config.get('extra_apt_packages') + ), + 'extra_python_packages': _format_pip_packages_for_shell( + config.get('extra_python_packages') + ), } -def lithops_pip_spec_from_config(config=None, default='lithops'): +def lithops_pip_spec_from_config(config=None, default: str = 'lithops') -> str: """ - Build a minimal pip spec from lithops config (avoid lithops[all] on VMs). - Standalone master/workers always need the redis extra for the job queue. + Returns the pip spec the VMs install, holding only the extras the + configured backends need. Installing lithops[all] on a VM would pull in + every cloud SDK, and the redis extra is always needed for the job queue """ if not config: return default @@ -186,7 +222,7 @@ def lithops_pip_spec_from_config(config=None, default='lithops'): name = (config.get(key) or lithops_cfg.get(key) or '').lower() if name.startswith('gcp'): extras.add('gcp') - elif name.startswith('aws') or name in ('aws_s3', 'aws_sqs'): + elif name.startswith('aws'): extras.add('aws') elif name.startswith('azure'): extras.add('azure') @@ -204,16 +240,18 @@ def lithops_pip_spec_from_config(config=None, default='lithops'): def get_host_setup_script( - docker=True, - run_install=True, - lithops_pip_spec='lithops', - extra_apt_packages='', - extra_python_packages='', -): + docker: bool = True, + run_install: bool = True, + lithops_pip_spec: str = 'lithops', + extra_apt_packages: str = '', + extra_python_packages: str = '', +) -> str: """ - Returns the script necessary for installing a lithops VM host. - Set run_install=False when appending master/worker setup (they run install first). - extra_apt_packages/extra_python_packages are pre-validated space-separated strings. + Returns the script that installs everything a Lithops VM host needs. + + Pass run_install=False when the master or worker setup is appended to it, + as those run the installation themselves. The extra package arguments are + space separated strings that have already been validated """ script = f"""#!/bin/bash mkdir -p {SA_INSTALL_DIR}; @@ -255,7 +293,9 @@ def get_host_setup_script( set -e export DEBIAN_FRONTEND=noninteractive export DOCKER_REQUIRED={str(docker).lower()}; - command -v docker >/dev/null 2>&1 || {{ export INSTALL_DOCKER=true; export INSTALL_LITHOPS_DEPS=true;}}; + command -v docker >/dev/null 2>&1 || {{ + export INSTALL_DOCKER=true; export INSTALL_LITHOPS_DEPS=true; + }}; command -v unzip >/dev/null 2>&1 || {{ export INSTALL_LITHOPS_DEPS=true; }}; command -v pip3 >/dev/null 2>&1 || {{ export INSTALL_LITHOPS_DEPS=true; }}; @@ -264,7 +304,8 @@ def get_host_setup_script( echo "--> Installing Docker repository" apt_install update apt_install install -y apt-transport-https ca-certificates curl gnupg software-properties-common - curl -fsSL https://download.docker.com/linux/ubuntu/gpg | gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg + curl -fsSL https://download.docker.com/linux/ubuntu/gpg | + gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg DOCKER_APT="deb [arch=amd64 signed-by=/usr/share/keyrings/docker-archive-keyring.gpg]" DOCKER_APT="$DOCKER_APT https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable" echo "$DOCKER_APT" > /etc/apt/sources.list.d/docker.list @@ -318,20 +359,30 @@ def get_host_setup_script( return script -def docker_login(config): +def docker_login(config) -> str: + """ + Returns the script line that logs into a private container registry, or + an empty string when no credentials are configured + """ backend = config['backend'] if all(k in config[backend] for k in ("docker_server", "docker_user", "docker_password")): user = config[backend]['docker_user'] passwd = config[backend]['docker_password'] server = config[backend]['docker_server'] - return f"""docker login -u {user} -p {passwd} {server} >> /tmp/kuku 2>&1 + login = ( + f"printf '%s' {shlex.quote(passwd)} | docker login " + f"-u {shlex.quote(user)} --password-stdin {shlex.quote(server)}" + ) + return f"""{login} >> {SA_SETUP_LOG_FILE} 2>&1 """ return "" -def get_master_setup_script(config, vm_data): +def get_master_setup_script(config, vm_data) -> str: """ - Returns master VM installation script + Returns the script that turns a VM into the Lithops master: it unpacks the + package, starts the master service, and generates the key pair the master + uses to reach the workers it creates """ script = docker_login(config) script += f""" @@ -360,7 +411,8 @@ def get_master_setup_script(config, vm_data): ssh-keygen -f $USER_HOME/.ssh/lithops_id_rsa -t rsa -N ''; cp $USER_HOME/.ssh/lithops_id_rsa $USER_HOME/.ssh/id_rsa cp $USER_HOME/.ssh/lithops_id_rsa.pub $USER_HOME/.ssh/id_rsa.pub - chown ${{SUDO_USER}}:${{SUDO_USER}} $USER_HOME/.ssh/lithops_id_rsa* $USER_HOME/.ssh/id_rsa $USER_HOME/.ssh/id_rsa.pub + chown ${{SUDO_USER}}:${{SUDO_USER}} $USER_HOME/.ssh/lithops_id_rsa* + chown ${{SUDO_USER}}:${{SUDO_USER}} $USER_HOME/.ssh/id_rsa $USER_HOME/.ssh/id_rsa.pub chmod 600 $USER_HOME/.ssh/lithops_id_rsa $USER_HOME/.ssh/id_rsa chmod 644 $USER_HOME/.ssh/lithops_id_rsa.pub $USER_HOME/.ssh/id_rsa.pub cp $USER_HOME/.ssh/lithops_id_rsa /root/.ssh/lithops_id_rsa @@ -378,24 +430,42 @@ def get_master_setup_script(config, vm_data): return script -def get_worker_setup_script(config, vm_data): +def _worker_service_commands(config: Dict[str, Any]) -> Tuple[str, str, str]: + """ + Returns the systemd ExecStartPre, ExecStart and ExecStop of the worker + service. A container runtime runs the worker inside the image, so it has + to remove a leftover container before starting and after stopping + """ + if not is_container_runtime(config['runtime']): + identity = 'id' + start = f"/usr/bin/python3 {SA_INSTALL_DIR}/worker.py" + return identity, start, identity + + gpu = '--gpus all ' if config.get('use_gpu') else '' + uid = os.getuid() + gid = os.getgid() + user = os.getenv('USER', 'root') + runtime = config['runtime'] + rm = '-docker rm -f lithops_worker' + start = ( + 'docker run --rm --name lithops_worker ' + f'{gpu}' + f'--user {uid}:{gid} ' + f'--env USER={user} --env DOCKER=Lithops ' + f'-p {SA_WORKER_SERVICE_PORT}:{SA_WORKER_SERVICE_PORT} ' + f'-v {SA_INSTALL_DIR}:{SA_INSTALL_DIR} -v /tmp:/tmp ' + f'--entrypoint "python3" {runtime} {SA_INSTALL_DIR}/worker.py' + ) + return rm, start, rm + + +def get_worker_setup_script(config, vm_data) -> str: """ - Returns worker VM installation script - this script is expected to be executed only from Master VM + Returns the script that turns a VM into a Lithops worker, which only the + master runs, as it is the one holding the key the worker has to trust """ - if config['runtime'].startswith(('python', '/')): - cmd_pre = cmd_stop = "id" - cmd_start = f"/usr/bin/python3 {SA_INSTALL_DIR}/worker.py" - else: - cmd_pre = '-docker rm -f lithops_worker' - cmd_start = 'docker run --rm --name lithops_worker ' - cmd_start += '--gpus all ' if config["use_gpu"] else '' - cmd_start += f'--user {os.getuid()}:{os.getgid()} ' - cmd_start += f'--env USER={os.getenv("USER", "root")} --env DOCKER=Lithops ' - cmd_start += f'-p {SA_WORKER_SERVICE_PORT}:{SA_WORKER_SERVICE_PORT} ' - cmd_start += f'-v {SA_INSTALL_DIR}:{SA_INSTALL_DIR} -v /tmp:/tmp ' - cmd_start += f'--entrypoint "python3" {config["runtime"]} {SA_INSTALL_DIR}/worker.py' - cmd_stop = '-docker rm -f lithops_worker' + cmd_pre, cmd_start, cmd_stop = _worker_service_commands(config) + unit_file = WORKER_SERVICE_FILE.format(cmd_pre, cmd_start, cmd_stop) script = docker_login(config) script += f""" @@ -407,7 +477,7 @@ def get_worker_setup_script(config, vm_data): }} USER_HOME=$(eval echo ~${{SUDO_USER}}); setup_service(){{ - echo '{WORKER_SERVICE_FILE.format(cmd_pre, cmd_start, cmd_stop)}' > /etc/systemd/system/{WORKER_SERVICE_NAME}; + echo '{unit_file}' > /etc/systemd/system/{WORKER_SERVICE_NAME}; chmod 644 /etc/systemd/system/{WORKER_SERVICE_NAME}; systemctl daemon-reload; systemctl stop {WORKER_SERVICE_NAME}; @@ -424,11 +494,17 @@ def get_worker_setup_script(config, vm_data): if "ssh_credentials" in vm_data: ssh_user = vm_data['ssh_credentials']['username'] home_dir = '/root' if ssh_user == 'root' else f'/home/{ssh_user}' + master_pub_key = '' try: - master_pub_key = open(f'{home_dir}/.ssh/lithops_id_rsa.pub', 'r').read() - except Exception: - master_pub_key = '' - script += f""" + with open(f'{home_dir}/.ssh/lithops_id_rsa.pub', 'r') as key_file: + master_pub_key = key_file.read() + except OSError: + # The master generates this key on its own setup, so a worker + # created before that has nothing to authorize yet + pass + + if master_pub_key: + script += f""" if ! grep -qF "{master_pub_key}" "$USER_HOME/.ssh/authorized_keys"; then echo "{master_pub_key}" >> $USER_HOME/.ssh/authorized_keys; fi diff --git a/lithops/standalone/worker.py b/lithops/standalone/worker.py index 03cb77faa..01b6377b8 100644 --- a/lithops/standalone/worker.py +++ b/lithops/standalone/worker.py @@ -17,18 +17,20 @@ import os import json +import time import redis import flask import logging import signal import subprocess as sp +from contextlib import contextmanager from pathlib import Path from threading import Thread from functools import partial from gevent.pywsgi import WSGIServer from concurrent.futures import ThreadPoolExecutor -from lithops.utils import setup_lithops_logger +from lithops.utils import setup_lithops_logger, log_prefix from lithops.standalone.keeper import BudgetKeeper from lithops.standalone.utils import JobStatus, StandaloneMode, WorkerStatus from lithops.constants import ( @@ -41,19 +43,23 @@ LOGS_DIR, SA_CONFIG_FILE, SA_WORKER_DATA_FILE, - SA_WORKER_SERVICE_PORT + SA_WORKER_SERVICE_PORT, ) -os.makedirs(LITHOPS_TEMP_DIR, exist_ok=True) -os.makedirs(JOBS_DIR, exist_ok=True) -os.makedirs(LOGS_DIR, exist_ok=True) - -log_format = "%(asctime)s\t[%(levelname)s] %(name)s:%(lineno)s -- %(message)s" -setup_lithops_logger(logging.DEBUG, filename=SA_WORKER_LOG_FILE, log_format=log_format) logger = logging.getLogger('lithops.standalone.worker') app = flask.Flask(__name__) +_LOG_FORMAT = ( + "%(asctime)s\t[%(levelname)s] %(name)s:%(lineno)s -- %(message)s" +) +# Reuse/consume workers sit on BRPOP while they wait for the next job. A +# timeout of 0 holds that connection forever, and a connection that goes +# stale (idle NAT, Redis keepalive) never sees the LPUSH of the next job: +# /ping still reports the process as free, the job sits in the queue, and +# the budget keeper eventually dismantles the worker as idle. +_QUEUE_POLL_TIMEOUT = 5 + redis_client = None budget_keeper = None @@ -62,10 +68,40 @@ canceled = [] +def _configure_logging(): + """Creates the directories the worker writes to and opens its log""" + os.makedirs(LITHOPS_TEMP_DIR, exist_ok=True) + os.makedirs(JOBS_DIR, exist_ok=True) + os.makedirs(LOGS_DIR, exist_ok=True) + setup_lithops_logger( + logging.DEBUG, filename=SA_WORKER_LOG_FILE, log_format=_LOG_FORMAT + ) + + +def _kill_process_group(process): + """ + Kills a task process along with everything it forked, which is why the + whole process group goes down and not just the process itself + """ + if not process or process.poll() is not None: + return + try: + os.killpg(os.getpgid(process.pid), signal.SIGKILL) + except Exception: + pass + + @app.route('/ping', methods=['GET']) def ping(): - idle_count = sum(1 for worker in worker_threads.values() if worker['status'] == WorkerStatus.IDLE.value) - busy_count = sum(1 for worker in worker_threads.values() if worker['status'] == WorkerStatus.BUSY.value) + """Reports how many of the worker processes are busy and how many free""" + idle_count = sum( + 1 for worker in worker_threads.values() + if worker['status'] == WorkerStatus.IDLE.value + ) + busy_count = sum( + 1 for worker in worker_threads.values() + if worker['status'] == WorkerStatus.BUSY.value + ) response = flask.jsonify({'busy': busy_count, 'free': idle_count}) response.status_code = 200 return response @@ -73,127 +109,185 @@ def ping(): @app.route('/ttd', methods=['GET']) def ttd(): + """Reports the seconds left before this worker stops itself""" if budget_keeper: - ttd = budget_keeper.get_time_to_dismantle() + ttd_value = budget_keeper.get_time_to_dismantle() else: - ttd = "Disabled" - return str(ttd), 200 + ttd_value = "Disabled" + return str(ttd_value), 200 @app.route('/stop/', methods=['POST']) def stop(job_key): + """Kills the task processes of a job and marks its tasks as done""" logger.debug(f'Received SIGTERM: Stopping job process {job_key}') canceled.append(job_key) - for job_key_call_id in job_processes: - if job_key_call_id.startswith(job_key): - PID = job_processes[job_key_call_id].pid - PGID = os.getpgid(PID) - logger.debug(f"Killing Job {job_key} - PID {PID}") - os.killpg(PGID, signal.SIGKILL) - Path(os.path.join(JOBS_DIR, job_key_call_id + '.done')).touch() - job_processes[job_key_call_id] = None + # A snapshot, because the consumer threads add and remove entries while + # this runs + for job_key_call_id, process in list(job_processes.items()): + if not job_key_call_id.startswith(job_key): + continue + logger.debug(f"Killing Job {job_key} - PID {getattr(process, 'pid', None)}") + _kill_process_group(process) + Path(os.path.join(JOBS_DIR, job_key_call_id + '.done')).touch() + job_processes.pop(job_key_call_id, None) response = flask.jsonify({'response': 'cancel'}) response.status_code = 200 return response -def notify_worker_active(worker_name): +@contextmanager +def _reported(what): + """ + Runs a piece of Redis bookkeeping, reporting a failure instead of raising + it: a worker that cannot tell the master what it is doing still has to run + the tasks it was given + """ try: - redis_client.hset(f"worker:{worker_name}", 'status', WorkerStatus.ACTIVE.value) + yield except Exception as e: - logger.error(e) + logger.error(f'Could not {what}: {e}') + + +def notify_worker_active(worker_name): + """Tells the master this worker is up""" + with _reported(f'mark worker {worker_name} as active'): + redis_client.hset( + f"worker:{worker_name}", 'status', WorkerStatus.ACTIVE.value + ) def notify_worker_idle(worker_name): - try: - data = {'status': WorkerStatus.IDLE.value, 'runtime': '', 'worker_processes': ''} - redis_client.hset(f"worker:{worker_name}", mapping=data) - except Exception as e: - logger.error(e) + """Tells the master this worker is free to take another job""" + with _reported(f'mark worker {worker_name} as idle'): + redis_client.hset(f"worker:{worker_name}", mapping={ + 'status': WorkerStatus.IDLE.value, + 'runtime': '', + 'worker_processes': '', + }) def notify_worker_stop(worker_name): - try: - redis_client.hset(f"worker:{worker_name}", 'status', WorkerStatus.STOPPED.value) - except Exception as e: - logger.error(e) + """Tells the master this worker is stopping""" + with _reported(f'mark worker {worker_name} as stopped'): + redis_client.hset( + f"worker:{worker_name}", 'status', WorkerStatus.STOPPED.value + ) def notify_worker_delete(worker_name): - try: + """Tells the master this worker is going away for good""" + with _reported(f'delete worker {worker_name}'): redis_client.delete(f"worker:{worker_name}") - except Exception as e: - logger.error(e) def notify_task_start(job_key, call_id): - try: + """Marks the job as running, the first time one of its tasks starts""" + with _reported(f'mark job {job_key} as running'): if redis_client.hget(f"job:{job_key}", 'status') == JobStatus.SUBMITTED.value: - redis_client.hset(f"job:{job_key}", 'status', JobStatus.RUNNING.value) - except Exception as e: - logger.error(e) + redis_client.hset( + f"job:{job_key}", 'status', JobStatus.RUNNING.value + ) def notify_task_done(job_key, call_id): - try: + """ + Counts a finished task, and marks the job as done once every one of its + tasks has been counted + """ + with _reported(f'mark task {call_id} of job {job_key} as done'): done_tasks = int(redis_client.rpush(f"tasksdone:{job_key}", call_id)) if int(redis_client.hget(f"job:{job_key}", 'total_tasks')) == done_tasks: redis_client.hset(f"job:{job_key}", 'status', JobStatus.DONE.value) - except Exception as e: - logger.error(e) + + +def _wait_for_task(work_queue_name, exec_mode): + """ + Returns the next task payload from the work queue, or None when there is + nothing to run right now. Create mode treats that as the end of the job. + The other modes poll with a timeout so that a stale BRPOP is dropped and + opened again, instead of waiting forever on a connection that will never + see the next job + """ + if exec_mode == StandaloneMode.CREATE.value: + return redis_client.rpop(work_queue_name) + + item = redis_client.brpop(work_queue_name, timeout=_QUEUE_POLL_TIMEOUT) + if item is None: + return None + _key, task_payload_str = item + return task_payload_str def redis_queue_consumer(pid, work_queue_name, exec_mode, backend): + """ + Takes tasks from the work queue and runs them one after another, until the + queue runs dry in create mode, or forever in the modes where the worker + waits for the jobs still to come + """ worker_threads[pid]['status'] = WorkerStatus.IDLE.value - logger.info(f"Redis consumer process {pid} started") while True: - if exec_mode == StandaloneMode.CREATE.value: - task_payload_str = redis_client.rpop(work_queue_name) - if task_payload_str is None: + try: + task_payload_str = _wait_for_task(work_queue_name, exec_mode) + except Exception as e: + logger.error( + f'Redis consumer {pid} could not read the queue: {e}' + ) + time.sleep(1) + continue + + if task_payload_str is None: + if exec_mode == StandaloneMode.CREATE.value: break - else: - key, task_payload_str = redis_client.brpop(work_queue_name) + continue worker_threads[pid]['status'] = WorkerStatus.BUSY.value - - task_payload = json.loads(task_payload_str) - - executor_id = task_payload['executor_id'] - job_id = task_payload['job_id'] - job_key = task_payload['job_key'] - call_id = task_payload['call_ids'][0] - job_key_call_id = f'{job_key}-{call_id}' - try: - logger.debug(f'ExecutorID {executor_id} | JobID {job_id} - Running ' - f'CallID {call_id} in the local worker (consumer {pid})') + task_payload = json.loads(task_payload_str) + executor_id = task_payload['executor_id'] + job_id = task_payload['job_id'] + job_key = task_payload['job_key'] + call_id = task_payload['call_ids'][0] + job_key_call_id = f'{job_key}-{call_id}' + + logger.debug( + f'{log_prefix(executor_id, job_id)} - Running ' + f'CallID {call_id} in the local worker (consumer {pid})' + ) notify_task_start(job_key, call_id) if budget_keeper: budget_keeper.add_job(job_key_call_id) task_filename = os.path.join(JOBS_DIR, f'{job_key_call_id}.task') - - with open(task_filename, 'w') as jl: - json.dump(task_payload, jl, default=str) - - cmd = ["python3", f"{SA_INSTALL_DIR}/runner.py", backend, task_filename] - log = open(RN_LOG_FILE, 'a') - process = sp.Popen(cmd, stdout=log, stderr=log, start_new_session=True) - job_processes[job_key_call_id] = process - process.communicate() # blocks until the process finishes - del job_processes[job_key_call_id] + with open(task_filename, 'w') as task_file: + json.dump(task_payload, task_file, default=str) + + cmd = [ + "python3", + f"{SA_INSTALL_DIR}/runner.py", + backend, + task_filename, + ] + with open(RN_LOG_FILE, 'a') as log: + process = sp.Popen( + cmd, stdout=log, stderr=log, start_new_session=True + ) + job_processes[job_key_call_id] = process + process.communicate() + # Popped, not deleted: the stop route may have taken it already + job_processes.pop(job_key_call_id, None) if os.path.exists(task_filename): os.remove(task_filename) Path(os.path.join(JOBS_DIR, f'{job_key_call_id}.done')).touch() - msg = f'ExecutorID {executor_id} | JobID {job_id} - ' + msg = f'{log_prefix(executor_id, job_id)} - ' if job_key in canceled: msg += f'CallID {call_id} execution canceled' else: @@ -209,58 +303,72 @@ def redis_queue_consumer(pid, work_queue_name, exec_mode, backend): def run_worker(): + """ + Entry point of the worker service: connects to the master, starts the + countdown that stops the instance, serves the control endpoints, and runs + one consumer per worker process until there is nothing left to run + """ global redis_client global budget_keeper - os.makedirs(LITHOPS_TEMP_DIR, exist_ok=True) - - # read the Lithops standaole configuration file - with open(SA_CONFIG_FILE, 'r') as cf: - standalone_config = json.load(cf) + _configure_logging() - # Read the VM data file that contains the instance id, the master IP, - # and the queue for getting tasks - with open(SA_WORKER_DATA_FILE, 'r') as ad: - worker_data = json.load(ad) + with open(SA_CONFIG_FILE, 'r') as config_file: + standalone_config = json.load(config_file) - # Start the redis client - redis_client = redis.Redis(host=worker_data['master_ip'], decode_responses=True) + with open(SA_WORKER_DATA_FILE, 'r') as data_file: + worker_data = json.load(data_file) - # Set the worker as Active + redis_client = redis.Redis( + host=worker_data['master_ip'], + decode_responses=True, + socket_keepalive=True, + ) notify_worker_active(worker_data['name']) - # Start the budget keeper. It is responsible to automatically terminate the - # worker after X seconds if worker_data['master_ip'] != worker_data['private_ip']: stop_callback = partial(notify_worker_stop, worker_data['name']) delete_callback = partial(notify_worker_delete, worker_data['name']) - budget_keeper = BudgetKeeper(standalone_config, worker_data, stop_callback, delete_callback) + budget_keeper = BudgetKeeper( + standalone_config, worker_data, stop_callback, delete_callback + ) budget_keeper.start() - # Start the http server. This will be used by the master VM to pìng this - # worker and for canceling tasks def run_wsgi(): - ip_address = "0.0.0.0" if os.getenv("DOCKER") == "Lithops" else worker_data['private_ip'] - server = WSGIServer((ip_address, SA_WORKER_SERVICE_PORT), app, log=app.logger) + """Serves the control endpoints of this worker""" + ip_address = ( + "0.0.0.0" if os.getenv("DOCKER") == "Lithops" + else worker_data['private_ip'] + ) + server = WSGIServer( + (ip_address, SA_WORKER_SERVICE_PORT), app, log=app.logger + ) server.serve_forever() + Thread(target=run_wsgi, daemon=True).start() - # Start the consumer threads - worker_processes = standalone_config[standalone_config['backend']]['worker_processes'] - worker_processes = CPU_COUNT if worker_processes == 'AUTO' else worker_processes - logger.info(f"Starting Worker - Instance type: {worker_data['instance_type']} - Runtime " - f"name: {standalone_config['runtime']} - Worker processes: {worker_processes}") + worker_processes = standalone_config[ + standalone_config['backend'] + ]['worker_processes'] + worker_processes = ( + CPU_COUNT if worker_processes == 'AUTO' else worker_processes + ) + logger.info( + f"Starting Worker - Instance type: {worker_data['instance_type']} - " + f"Runtime name: {standalone_config['runtime']} - " + f"Worker processes: {worker_processes}" + ) - # Create a ThreadPoolExecutor for cosnuming tasks redis_queue_consumer_futures = [] with ThreadPoolExecutor(max_workers=worker_processes) as executor: for i in range(worker_processes): worker_threads[i] = {} future = executor.submit( - redis_queue_consumer, i, + redis_queue_consumer, + i, worker_data['work_queue_name'], standalone_config['exec_mode'], - standalone_config['backend'] + standalone_config['backend'], ) redis_queue_consumer_futures.append(future) worker_threads[i]['future'] = future @@ -268,20 +376,16 @@ def run_wsgi(): for future in redis_queue_consumer_futures: future.result() - # Set the worker as idle if standalone_config['exec_mode'] == StandaloneMode.CONSUME.value: notify_worker_idle(worker_data['name']) - # run_worker will run forever in reuse mode. In create and consume mode it will - # run until there are no more tasks in the queue. logger.debug('Worker service finished') - try: - # Try to stop the current worker VM once no more pending tasks to run - # in case of create mode - budget_keeper.stop_instance() - except Exception: - pass + if budget_keeper: + try: + budget_keeper.stop_instance() + except Exception as e: + logger.error(f'Could not stop the instance: {e}') if __name__ == '__main__': diff --git a/lithops/storage/cloud_proxy.py b/lithops/storage/cloud_proxy.py index d3f0005c7..817e8b6be 100644 --- a/lithops/storage/cloud_proxy.py +++ b/lithops/storage/cloud_proxy.py @@ -17,15 +17,29 @@ import io import os as base_os from functools import partial +from typing import Any, Dict, Iterable, List, Optional, Union + from lithops.storage import Storage from lithops.utils import is_lithops_worker -from lithops.config import default_storage_config, load_yaml_config, extract_storage_config -from lithops.constants import JOBS_PREFIX, TEMP_PREFIX, LOGS_PREFIX, RUNTIMES_PREFIX +from lithops.config import ( + default_storage_config, + load_yaml_config, + extract_storage_config, +) +from lithops.constants import ( + JOBS_PREFIX, TEMP_PREFIX, LOGS_PREFIX, RUNTIMES_PREFIX, +) + + +_LITHOPS_PREFIXES = (JOBS_PREFIX, TEMP_PREFIX, LOGS_PREFIX, RUNTIMES_PREFIX) -def remove_lithops_keys(keys): - return list(filter(lambda key: not any([key.startswith(prefix) for prefix in [ - JOBS_PREFIX, TEMP_PREFIX, LOGS_PREFIX, RUNTIMES_PREFIX]]), keys)) +def remove_lithops_keys(keys: Iterable[str]) -> List[str]: + """ + Drops the keys Lithops itself writes, so that the proxy only shows the + data of the user + """ + return [key for key in keys if not key.startswith(_LITHOPS_PREFIXES)] # @@ -33,15 +47,20 @@ def remove_lithops_keys(keys): # class CloudStorage(Storage): - def __init__(self, config=None): + """ + Storage client that can be pickled, so that it can travel to a worker. + It keeps the configuration it was built from and builds a new client on + the other side, as the underlying backend clients are not picklable + """ + + def __init__(self, config: Optional[Union[str, Dict[str, Any]]] = None): if isinstance(config, str): config = load_yaml_config(config) self._config = extract_storage_config(config) + elif isinstance(config, dict) and 'lithops' in config: + self._config = extract_storage_config(config) elif isinstance(config, dict): - if 'lithops' in config: - self._config = extract_storage_config(config) - else: - self._config = config + self._config = config else: self._config = extract_storage_config(default_storage_config()) super().__init__(storage_config=self._config) @@ -53,20 +72,29 @@ def __setstate__(self, state): self.__init__(state) def put_data(self, key, data): + """Writes an object in the configured bucket""" return self.put_object(self.bucket, key, data) def get_data(self, key): + """Reads an object from the configured bucket""" return self.get_object(self.bucket, key) def delete_data(self, key): + """Deletes an object from the configured bucket""" self.delete_object(self.bucket, key) def list_bucket_keys(self, prefix=None): + """Lists the keys of the configured bucket""" return self.list_keys(self.bucket, prefix) class CloudFileProxy: - def __init__(self, cloud_storage=None): + """ + Stand-in for the os module that reads and writes objects in storage + instead of files. Anything it does not implement is served by os itself + """ + + def __init__(self, cloud_storage: Optional[CloudStorage] = None): self._storage = cloud_storage or CloudStorage() self.path = _path(self._storage) @@ -75,32 +103,48 @@ def __getattr__(self, name): return getattr(base_os, name) def open(self, filename, mode='r'): + """Opens an object as a file-like buffer""" return cloud_open(filename, mode=mode, cloud_storage=self._storage) def listdir(self, path='', suffix_dirs=False): + """ + Lists the names directly under a path, as os.listdir does. Keys are + flat in storage, so a name is the first segment left after the prefix, + and the ones that stand for a directory can be marked with a slash + """ if path == '': - prefix = '/' + prefix = '' elif path.startswith('/'): prefix = path[1:] else: prefix = path if path.endswith('/') else path + '/' - paths = self._storage.list_bucket_keys(prefix=prefix) names = set() - for p in paths: - if any([p.startswith(prefix) for prefix in [ - JOBS_PREFIX, TEMP_PREFIX, LOGS_PREFIX, RUNTIMES_PREFIX]]): - continue + for p in remove_lithops_keys(self._storage.list_bucket_keys(prefix=prefix)): p = p[len(prefix):] if p.startswith(prefix) else p if p.startswith('/'): p = p[1:] splits = p.split('/') - name = splits[0] + \ - '/' if suffix_dirs and len(splits) > 1 else splits[0] - names |= {name} + name = ( + splits[0] + '/' if suffix_dirs and len(splits) > 1 + else splits[0] + ) + names.add(name) return list(names) + def _walk_children(self, top, dirs, topdown, onerror, followlinks): + """Walks each subdirectory of a path, in the requested order""" + for dir_name in dirs: + yield from self.walk( + base_os.path.join(top, dir_name), + topdown, onerror, followlinks, + ) + def walk(self, top, topdown=True, onerror=None, followlinks=False): + """ + Walks a path yielding (top, dirs, files), as os.walk does, and yields + nothing at all when the path holds no key + """ dirs = [] files = [] @@ -111,77 +155,76 @@ def walk(self, top, topdown=True, onerror=None, followlinks=False): files.append(path) if dirs == [] and files == [] and not self.path.exists(top): - raise StopIteration - elif topdown: + return + if topdown: yield top, dirs, files - for dir_name in dirs: - for result in self.walk( - base_os.path.join( - top, - dir_name), - topdown, - onerror, - followlinks): - yield result + yield from self._walk_children( + top, dirs, topdown, onerror, followlinks + ) else: - for dir_name in dirs: - for result in self.walk( - base_os.path.join( - top, - dir_name), - topdown, - onerror, - followlinks): - yield result + yield from self._walk_children( + top, dirs, topdown, onerror, followlinks + ) yield top, dirs, files def remove(self, path): + """Deletes the object a path names""" self._storage.delete_data(path) def mkdir(self, *args, **kwargs): + """Does nothing: storage has no directories to create""" pass def makedirs(self, *args, **kwargs): + """Does nothing: storage has no directories to create""" pass class _path: - def __init__(self, cloud_storage=None): + """ + Stand-in for os.path that answers from the keys in the bucket. Anything + it does not implement is served by os.path itself + """ + + def __init__(self, cloud_storage: Optional[CloudStorage] = None): self._storage = cloud_storage or CloudStorage() def __getattr__(self, name): # we only reach here if the attr is not defined return getattr(base_os.path, name) - def isfile(self, path): - prefix = path - if path.startswith('/'): - prefix = path[1:] + def _prefix(self, path, as_dir=False): + """ + Turns a path into the key prefix that matches it, with a trailing + slash when only the contents of a directory should match + """ + prefix = path[1:] if path.startswith('/') else path + if as_dir and prefix != '' and not prefix.endswith('/'): + prefix = prefix + '/' + return prefix + def isfile(self, path): + """True when the path names one object and not a prefix of others""" + prefix = self._prefix(path) keys = remove_lithops_keys( - self._storage.list_bucket_keys( - prefix=prefix)) + self._storage.list_bucket_keys(prefix=prefix) + ) if len(keys) == 1: key = keys.pop() key = key[len(prefix):] return key == '' - else: - return False + return False def isdir(self, path): - prefix = path - if path.startswith('/'): - prefix = path[1:] - - if prefix != '' and not prefix.endswith('/'): - prefix = prefix + '/' - + """True when there is at least one object under the path""" + prefix = self._prefix(path, as_dir=True) keys = remove_lithops_keys( - self._storage.list_bucket_keys( - prefix=prefix)) + self._storage.list_bucket_keys(prefix=prefix) + ) return bool(keys) def exists(self, path): + """True when the path names an object or a directory holding one""" dirpath = path if path.endswith('/') else path + '/' for key in self._storage.list_bucket_keys(prefix=path): if key.startswith(dirpath) or key == path: @@ -189,41 +232,53 @@ def exists(self, path): return False -class DelayedBytesBuffer(io.BytesIO): +class _DelayedClose: + """ + Buffer that runs its action on close, which is what makes a write to + storage happen only once the caller is done writing + """ + + def close(self): + self._action(self.getvalue()) + super().close() + + +class DelayedBytesBuffer(_DelayedClose, io.BytesIO): + """Binary buffer that uploads what it holds when it is closed""" + def __init__(self, action, initial_bytes=None): super().__init__(initial_bytes) self._action = action - def close(self): - self._action(self.getvalue()) - io.BytesIO.close(self) +class DelayedStringBuffer(_DelayedClose, io.StringIO): + """Text buffer that uploads what it holds when it is closed""" -class DelayedStringBuffer(io.StringIO): def __init__(self, action, initial_value=None): super().__init__(initial_value) self._action = action - def close(self): - self._action(self.getvalue()) - io.StringIO.close(self) - def cloud_open(filename, mode='r', cloud_storage=None): + """ + Opens an object as a file-like buffer. Reading brings the whole object + into memory, and writing uploads it when the buffer is closed + """ storage = cloud_storage or CloudStorage() if 'r' in mode: + data = storage.get_data(filename) if 'b' in mode: # we could get_data(stream=True) but some streams are not seekable - return io.BytesIO(storage.get_data(filename)) - else: - return io.StringIO(storage.get_data(filename).decode()) + return io.BytesIO(data) + return io.StringIO(data.decode()) if 'w' in mode: action = partial(storage.put_data, filename) if 'b' in mode: return DelayedBytesBuffer(action) - else: - return DelayedStringBuffer(action) + return DelayedStringBuffer(action) + + raise ValueError(f"Unsupported mode '{mode}': only 'r' and 'w' are") if not is_lithops_worker(): diff --git a/lithops/storage/storage.py b/lithops/storage/storage.py index c14800643..c654d23e8 100644 --- a/lithops/storage/storage.py +++ b/lithops/storage/storage.py @@ -1,4 +1,4 @@ - +# # (C) Copyright IBM Corp. 2020 # (C) Copyright Cloudlab URV 2020 # @@ -20,7 +20,7 @@ import logging import itertools import importlib -from typing import Optional, List, Union, Dict, TextIO, BinaryIO, Any +from typing import Optional, List, Union, Dict, TextIO, BinaryIO, Any, Iterable from lithops.constants import CACHE_DIR, RUNTIMES_PREFIX, JOBS_PREFIX, TEMP_PREFIX from lithops.utils import is_lithops_worker @@ -32,16 +32,17 @@ RUNTIME_META_CACHE = {} COBJECTS_INDEX = itertools.count() +_INVALID_CO_BACKEND = "CloudObject: Invalid Storage backend" class Storage: """ - An Storage object is used by partitioner and other components to access - underlying storage backend without exposing the implementation details. + A Storage object is used by the partitioner and other components to access + the underlying storage backend without exposing the implementation details. """ def __init__(self, config=None, backend=None, storage_config=None): - """ Creates an Storage instance + """ Creates a Storage instance :param config: lithops configuration dict :param backend: storage backend name @@ -49,25 +50,28 @@ def __init__(self, config=None, backend=None, storage_config=None): :return: Storage instance. """ - if storage_config: self.config = storage_config else: - storage_config = default_storage_config( - config_data=config, backend=backend) - self.config = extract_storage_config(storage_config) + self.config = extract_storage_config( + default_storage_config(config_data=config, backend=backend) + ) self.backend = self.config['backend'] try: - module_location = f'lithops.storage.backends.{self.backend}' - sb_module = importlib.import_module(module_location) + sb_module = importlib.import_module( + f'lithops.storage.backends.{self.backend}' + ) StorageBackend = getattr(sb_module, 'StorageBackend') self.storage_handler = StorageBackend(self.config[self.backend]) - except Exception as e: - logger.error("An exception was produced trying to create the " - f"'{self.backend}' storage backend") - raise e + except Exception: + logger.error( + "There was an error trying to create the " + f"'{self.backend}' storage backend", + exc_info=True, + ) + raise bucket = self.config[self.backend].get('storage_bucket') self.bucket = bucket or self.storage_handler.generate_bucket_name() @@ -90,7 +94,8 @@ def get_storage_config(self) -> Dict: def create_bucket(self, bucket: str): """ - Creates a bucket if not exists. + Creates a bucket if it does not exist. Backends that create their + buckets on their own do nothing here. :param bucket: Name of the bucket """ @@ -108,74 +113,74 @@ def put_object(self, bucket: str, key: str, """ return self.storage_handler.put_object(bucket, key, body) - def get_object(self, - bucket: str, - key: str, - stream: Optional[bool] = False, - extra_get_args: Optional[Dict] = {}) -> Union[str, - bytes, - TextIO, - BinaryIO]: + def get_object(self, bucket: str, key: str, stream: Optional[bool] = False, + extra_get_args: Optional[Dict] = {}) -> Union[ + str, bytes, TextIO, BinaryIO]: """ Retrieves objects from the storage backend. :param bucket: Name of the bucket :param key: Key of the object :param stream: Get the object data or a file-like object - :param extra_get_args: Extra get arguments to be passed to the underlying backend implementation (dict). - For example, to specify the byte-range to read: ``extra_get_args={'Range': 'bytes=0-100'}``. + :param extra_get_args: Extra get arguments to be passed to the + underlying backend implementation (dict). For example, to specify + the byte-range to read: ``extra_get_args={'Range': 'bytes=0-100'}``. - :return: Object, as a binary array or as a file-like stream if parameter `stream` is enabled + :return: Object, as a binary array or as a file-like stream if + parameter `stream` is enabled """ return self.storage_handler.get_object( - bucket, key, stream, extra_get_args) + bucket, key, stream, extra_get_args + ) - def upload_file(self, - file_name: str, - bucket: str, + def upload_file(self, file_name: str, bucket: str, key: Optional[str] = None, extra_args: Optional[Dict] = {}, - config: Optional[Any] = None) -> Union[str, - bytes, - TextIO, - BinaryIO]: + config: Optional[Any] = None) -> Union[ + str, bytes, TextIO, BinaryIO]: """ - Upload a file to a bucket of the storage backend. (Multipart upload) + Uploads a file to a bucket of the storage backend. (Multipart upload) :param file_name: Name of the file to upload :param bucket: Name of the bucket :param key: Key of the object - :param extra_args: Extra get arguments to be passed to the underlying backend implementation (dict). - :param config: The transfer configuration to be used when performing the transfer (boto3.s3.transfer.TransferConfig). + :param extra_args: Extra get arguments to be passed to the underlying + backend implementation (dict). + :param config: The transfer configuration to be used when performing + the transfer (boto3.s3.transfer.TransferConfig). """ - return self.storage_handler.upload_file(file_name, bucket, key, extra_args, config) + return self.storage_handler.upload_file( + file_name, bucket, key, extra_args, config + ) - def download_file(self, - bucket: str, - key: str, + def download_file(self, bucket: str, key: str, file_name: Optional[str] = None, extra_args: Optional[Dict] = {}, - config: Optional[Any] = None) -> Union[str, - bytes, - TextIO, - BinaryIO]: + config: Optional[Any] = None) -> Union[ + str, bytes, TextIO, BinaryIO]: """ - Download a file from the storage backend. (Multipart download) + Downloads a file from the storage backend. (Multipart download) :param bucket: Name of the bucket :param key: Key of the object :param file_name: Name of the file to save the object data - :param extra_args: Extra get arguments to be passed to the underlying backend implementation (dict). - :param config: The transfer configuration to be used when performing the transfer (boto3.s3.transfer.TransferConfig). + :param extra_args: Extra get arguments to be passed to the underlying + backend implementation (dict). + :param config: The transfer configuration to be used when performing + the transfer (boto3.s3.transfer.TransferConfig). - :return: Object, as a binary array or as a file-like stream if parameter `stream` is enabled + :return: Object, as a binary array or as a file-like stream if + parameter `stream` is enabled """ - return self.storage_handler.download_file(bucket, key, file_name, extra_args, config) + return self.storage_handler.download_file( + bucket, key, file_name, extra_args, config + ) def head_object(self, bucket: str, key: str) -> Dict: """ - The HEAD operation retrieves metadata from an object without returning the object itself. This operation is - useful if you're only interested in an object's metadata. + The HEAD operation retrieves metadata from an object without returning + the object itself. This operation is useful if you're only interested + in an object's metadata. :param bucket: Name of the bucket :param key: Key of the object @@ -195,9 +200,10 @@ def delete_object(self, bucket: str, key: str): def delete_objects(self, bucket: str, key_list: List[str]): """ - This operation enables you to delete multiple objects from a bucket using a single HTTP request. - If you know the object keys that you want to delete, then this operation provides a suitable alternative - to sending individual delete requests, reducing per-request overhead. + This operation enables you to delete multiple objects from a bucket + using a single HTTP request. If you know the object keys that you want + to delete, then this operation provides a suitable alternative to + sending individual delete requests, reducing per-request overhead. :param bucket: Name of the bucket :param key_list: List of object keys @@ -206,9 +212,10 @@ def delete_objects(self, bucket: str, key_list: List[str]): def head_bucket(self, bucket: str) -> Dict: """ - This operation is useful to determine if a bucket exists and you have permission to access it. - The operation returns a 200 OK if the bucket exists and you have permission to access it. - Otherwise, the operation might return responses such as 404 Not Found and 403 Forbidden. + This operation is useful to determine if a bucket exists and you have + permission to access it. The operation returns a 200 OK if the bucket + exists and you have permission to access it. Otherwise, the operation + might return responses such as 404 Not Found and 403 Forbidden. :param bucket: Name of the bucket @@ -216,28 +223,27 @@ def head_bucket(self, bucket: str) -> Dict: """ return self.storage_handler.head_bucket(bucket) - def list_objects(self, - bucket: str, - prefix: Optional[str] = None, - match_pattern: Optional[str] = None) -> List[Dict[str, - Any]]: + def list_objects(self, bucket: str, prefix: Optional[str] = None, + match_pattern: Optional[str] = None) -> List[Dict[str, Any]]: """ - Returns all of the object keys in a bucket. For each object, the list contains a dictionary - with at least the object key ('Key') and the size in bytes ('Size'). Additional fields may be - present, depending on the backend implementation. + Returns all of the object keys in a bucket. For each object, the list + contains a dictionary with at least the object key ('Key') and the size + in bytes ('Size'). Additional fields may be present, depending on the + backend implementation. :param bucket: Name of the bucket :param prefix: Key prefix for filtering - :return: List of dictionaries containing at least 'Key' and 'Size' for each object + :return: List of dictionaries containing at least 'Key' and 'Size' + for each object """ - return self.storage_handler.list_objects(bucket, prefix, match_pattern) def list_keys(self, bucket, prefix=None) -> List[str]: """ - Similar to list_objects(), it returns all of the object keys in a bucket. - For each object, the list contains only the names of the objects (keys). + Similar to list_objects(), it returns all of the object keys in a + bucket. For each object, the list contains only the names of the + objects (keys). :param bucket: Name of the bucket :param prefix: Key prefix for filtering @@ -246,17 +252,20 @@ def list_keys(self, bucket, prefix=None) -> List[str]: """ return self.storage_handler.list_keys(bucket, prefix) - def put_cloudobject(self, - body: Union[str, - bytes, - TextIO, - BinaryIO], + def _cloudobject_location(self, cloudobject: utils.CloudObject): + """Returns the bucket and the key a CloudObject of this backend lives in""" + if cloudobject.backend != self.backend: + raise Exception(_INVALID_CO_BACKEND) + return cloudobject.bucket, cloudobject.key + + def put_cloudobject(self, body: Union[str, bytes, TextIO, BinaryIO], bucket: Optional[str] = None, key: Optional[str] = None) -> utils.CloudObject: """ - Put a CloudObject into storage. + Puts a CloudObject into storage. - :param body: Data content, can be a string or byte array or a text/bytes file-like object + :param body: Data content, can be a string or byte array or a + text/bytes file-like object :param bucket: Destination bucket :param key: Destination key @@ -264,83 +273,62 @@ def put_cloudobject(self, """ prefix = os.environ.get('__LITHOPS_SESSION_ID', '') coid = hex(next(COBJECTS_INDEX))[2:] - coname = 'cloudobject_{}'.format(coid) + coname = f'cloudobject_{coid}' name = '/'.join([prefix, coname]) if prefix else coname key = key or '/'.join([TEMP_PREFIX, name]) bucket = bucket or self.bucket self.storage_handler.put_object(bucket, key, body) - return utils.CloudObject(self.backend, bucket, key) - def get_cloudobject(self, - cloudobject: utils.CloudObject, - stream: Optional[bool] = False) -> Union[str, - bytes, - TextIO, - BinaryIO]: + def get_cloudobject(self, cloudobject: utils.CloudObject, + stream: Optional[bool] = False) -> Union[ + str, bytes, TextIO, BinaryIO]: """ - Get a CloudObject's content from storage. + Gets the content of a CloudObject from storage. :param cloudobject: CloudObject instance :param stream: Get the object data or a file-like object :return: Cloud object content """ - if cloudobject.backend == self.backend: - bucket = cloudobject.bucket - key = cloudobject.key - return self.storage_handler.get_object(bucket, key, stream=stream) - else: - raise Exception("CloudObject: Invalid Storage backend") + bucket, key = self._cloudobject_location(cloudobject) + return self.storage_handler.get_object(bucket, key, stream=stream) def delete_cloudobject(self, cloudobject: utils.CloudObject): """ - Delete a CloudObject from storage. + Deletes a CloudObject from storage. :param cloudobject: CloudObject instance """ - if cloudobject.backend == self.backend: - bucket = cloudobject.bucket - key = cloudobject.key - return self.storage_handler.delete_object(bucket, key) - else: - raise Exception("CloudObject: Invalid Storage backend") + bucket, key = self._cloudobject_location(cloudobject) + return self.storage_handler.delete_object(bucket, key) def delete_cloudobjects(self, cloudobjects: List[utils.CloudObject]): """ - Delete multiple CloudObjects from storage. + Deletes multiple CloudObjects from storage. :param cloudobjects: List of CloudObject instances """ - cobjs = {} + keys_per_bucket = {} for co in cloudobjects: - if co.backend not in cobjs: - cobjs[co.backend] = {} - if co.bucket not in cobjs[co.backend]: - cobjs[co.backend][co.bucket] = [] - cobjs[co.backend][co.bucket].append(co.key) + # Checked before deleting anything, so that a foreign object in + # the list does not leave the others half deleted + if co.backend != self.backend: + raise Exception(_INVALID_CO_BACKEND) + keys_per_bucket.setdefault(co.bucket, []).append(co.key) - for backend in cobjs: - if backend == self.backend: - for bucket in cobjs[backend]: - self.storage_handler.delete_objects( - bucket, cobjs[backend][co.bucket]) - else: - raise Exception("CloudObject: Invalid Storage backend") + for bucket, keys in keys_per_bucket.items(): + self.storage_handler.delete_objects(bucket, keys) class InternalStorage: """ - An InternalStorage object is used by executors and other components to access - underlying storage backend without exposing the the implementation details. + An InternalStorage object is used by executors and other components to + access the underlying storage backend without exposing the implementation + details. Every key it reads and writes lives in the configured bucket """ - def __init__(self, storage_config): - """ Creates an InternalStorage instance - :param storage_config: Storage config dictionary - - :return: InternalStorage instance - """ + def __init__(self, storage_config: Dict[str, Any]): self.storage = Storage(storage_config=storage_config) self.backend = self.storage.backend self.bucket = self.storage.bucket @@ -348,94 +336,82 @@ def __init__(self, storage_config): if not self.bucket: raise Exception( f"'storage_bucket' is mandatory under '{self.backend}'" - " section of the configuration") + " section of the configuration" + ) self.storage.create_bucket(self.bucket) def get_client(self): - """ - Retrieves the underlying storage client. - :return: storage backend client - """ + """Returns the client of the underlying storage backend""" return self.storage.get_client() def get_storage_config(self): - """ - Retrieves the configuration of this storage handler. - :return: storage configuration - """ + """Returns the configuration of this storage handler""" return self.storage.get_storage_config() def put_data(self, key, data): - """ - Put data object into storage. - :param key: data key - :param data: data content - :return: None - """ + """Writes the data of a job""" return self.storage.put_object(self.bucket, key, data) def put_func(self, key, func): - """ - Put serialized function into storage. - :param key: function key - :param func: serialized function - :return: None - """ + """Writes a serialized function""" return self.storage.put_object(self.bucket, key, func) def get_data(self, key, stream=False, extra_get_args={}): - """ - Get data object from storage. - :param key: data key - :return: data content - """ - return self.storage.get_object( - self.bucket, key, stream, extra_get_args) + """Reads the data of a job, as bytes or as a stream""" + return self.storage.get_object(self.bucket, key, stream, extra_get_args) def get_func(self, key): - """ - Get serialized function from storage. - :param key: function key - :return: serialized function - """ + """Reads a serialized function""" return self.storage.get_object(self.bucket, key) def del_data(self, key): - """ - Deletes data from storage. - :param key: data key - :return: None - """ + """Deletes the data of a job""" return self.storage.delete_object(self.bucket, key) - def get_job_status(self, executor_id): - """ - Get the status of a callset. - :param executor_id: executor's ID - :return: A list of call IDs that have updated status. + def get_job_status(self, executor_id, job_ids: Optional[Iterable[str]] = None): """ - callset_prefix = '/'.join([JOBS_PREFIX, executor_id]) - keys = self.storage.list_keys(self.bucket, callset_prefix) - - running_keys = [k.split('/') - for k in keys if utils.init_key_suffix in k] - running_callids = [(tuple(k[1].rsplit("-", 1) + [k[2]]), - k[3].replace(utils.init_key_suffix, '')) - for k in running_keys] + Returns the ids of the calls that have started and of the ones that + have finished, as two sets. - done_keys = [k.split('/')[1:] - for k in keys if utils.status_key_suffix in k] - done_callids = [tuple(k[0].rsplit("-", 1) + [k[1]]) for k in done_keys] + Listing the prefix of each given job keeps finished jobs out of the + listing; without job_ids the whole executor prefix is listed + """ + if job_ids: + keys = [] + for job_id in job_ids: + prefix = '/'.join([ + JOBS_PREFIX, utils.create_job_key(executor_id, job_id) + ]) + keys.extend(self.storage.list_keys(self.bucket, prefix)) + else: + callset_prefix = '/'.join([JOBS_PREFIX, executor_id]) + keys = self.storage.list_keys(self.bucket, callset_prefix) + + running_keys = [ + k.split('/') for k in keys if utils.init_key_suffix in k + ] + running_callids = [ + ( + tuple(k[1].rsplit("-", 1) + [k[2]]), + k[3].replace(utils.init_key_suffix, ''), + ) + for k in running_keys + ] + + done_keys = [ + k.split('/')[1:] for k in keys if utils.status_key_suffix in k + ] + done_callids = [ + tuple(k[0].rsplit("-", 1) + [k[1]]) for k in done_keys + ] return set(running_callids), set(done_callids) def get_call_status(self, executor_id, job_id, call_id): """ - Get status of a call. - :param executor_id: executor ID of the call - :param call_id: call ID of the call - :return: A dictionary containing call's status, or None if no updated status + Returns the status of a single call, or None while it has not been + written yet """ status_key = utils.create_status_key(executor_id, job_id, call_id) try: @@ -446,10 +422,8 @@ def get_call_status(self, executor_id, job_id, call_id): def get_call_output(self, executor_id, job_id, call_id): """ - Get the output of a call. - :param executor_id: executor ID of the call - :param call_id: call ID of the call - :return: Output of the call. + Returns the serialized result of a single call, or None while it has + not been written yet """ output_key = utils.create_output_key(executor_id, job_id, call_id) try: @@ -457,88 +431,107 @@ def get_call_output(self, executor_id, job_id, call_id): except utils.StorageNoSuchKeyError: return None - def get_runtime_meta(self, key): + def _runtime_meta_refs(self, key): """ - Get the metadata given a runtime name. - :param runtime: name of the runtime - :return: runtime metadata + Returns where the metadata of a runtime lives: the path parts of the + local cache file, the key of the in-memory cache, and the storage key, + which is posix even when the local path is not """ path = [RUNTIMES_PREFIX, key + ".meta.json"] - filename_local_path = os.path.join(CACHE_DIR, *path) + cache_key = '/'.join(path) + return path, cache_key, cache_key.replace('\\', '/') + + def _local_runtime_meta_path(self, key): + """Returns the path of the local disk cache file of a runtime""" + path, _, _ = self._runtime_meta_refs(key) + return os.path.join(CACHE_DIR, *path) - if '/'.join(path) in RUNTIME_META_CACHE: + def _write_runtime_meta_file(self, filename_local_path, runtime_meta): + """Writes the metadata of a runtime to the local disk cache""" + os.makedirs(os.path.dirname(filename_local_path), exist_ok=True) + with open(filename_local_path, "w") as f: + f.write(json.dumps(runtime_meta)) + + def _cached_runtime_meta(self, cache_key, filename_local_path): + """ + Returns the metadata of a runtime from the memory cache, or from the + disk cache, which a worker does not use because its disk is not the + one that wrote it + """ + if cache_key in RUNTIME_META_CACHE: logger.debug("Runtime metadata found in local memory cache") - return RUNTIME_META_CACHE['/'.join(path)] + return RUNTIME_META_CACHE[cache_key] - elif not is_lithops_worker() and os.path.exists(filename_local_path): - logger.debug("Runtime metadata found in local disk cache") - with open(filename_local_path, "r") as f: - runtime_meta = json.loads(f.read()) - RUNTIME_META_CACHE['/'.join(path)] = runtime_meta + if is_lithops_worker() or not os.path.exists(filename_local_path): + return None + + logger.debug("Runtime metadata found in local disk cache") + with open(filename_local_path, "r") as f: + runtime_meta = json.loads(f.read()) + RUNTIME_META_CACHE[cache_key] = runtime_meta + return runtime_meta + + def get_runtime_meta(self, key): + """ + Returns the metadata of a runtime, looking in the memory cache, then + the disk cache, then storage. Returns None when the runtime has no + metadata yet, which is what tells the caller to deploy it + """ + _, cache_key, obj_key = self._runtime_meta_refs(key) + filename_local_path = self._local_runtime_meta_path(key) + + runtime_meta = self._cached_runtime_meta(cache_key, filename_local_path) + if runtime_meta is not None: return runtime_meta - else: - logger.debug( - "Runtime metadata not found in local cache. Retrieving it from storage") - try: - obj_key = '/'.join(path).replace('\\', '/') - logger.debug( - 'Trying to download runtime metadata from: {}://{}/{}' .format( - self.backend, self.bucket, obj_key)) - json_str = self.storage.get_object(self.bucket, obj_key) - logger.debug('Runtime metadata found in storage') - runtime_meta = json.loads(json_str.decode("ascii")) - - # Save runtime meta to cache - try: - if not os.path.exists( - os.path.dirname(filename_local_path)): - os.makedirs(os.path.dirname(filename_local_path)) - - with open(filename_local_path, "w") as f: - f.write(json.dumps(runtime_meta)) - except Exception as e: - logger.error( - "Could not save runtime meta to local cache: {}".format(e)) - - RUNTIME_META_CACHE['/'.join(path)] = runtime_meta - return runtime_meta - except utils.StorageNoSuchKeyError: - logger.debug('Runtime metadata not found in storage') - return None + logger.debug( + "Runtime metadata not found in local cache. Retrieving it from storage" + ) + logger.debug( + 'Trying to download runtime metadata from: ' + f'{self.backend}://{self.bucket}/{obj_key}' + ) + try: + json_str = self.storage.get_object(self.bucket, obj_key) + except utils.StorageNoSuchKeyError: + logger.debug('Runtime metadata not found in storage') + return None + + logger.debug('Runtime metadata found in storage') + runtime_meta = json.loads(json_str.decode("ascii")) + + try: + self._write_runtime_meta_file(filename_local_path, runtime_meta) + except Exception as e: + # A cache that cannot be written only costs the next download + logger.error(f"Could not save runtime meta to local cache: {e}") + + RUNTIME_META_CACHE[cache_key] = runtime_meta + return runtime_meta def put_runtime_meta(self, key, runtime_meta): """ - Put the metadata given a runtime config. - :param runtime: name of the runtime - :param runtime_meta metadata + Writes the metadata of a runtime to storage, and to the local disk + cache unless this is a worker, whose disk nothing else reads """ - path = [RUNTIMES_PREFIX, key + ".meta.json"] - obj_key = '/'.join(path).replace('\\', '/') - logger.debug("Uploading runtime metadata to: {}://{}/{}" - .format(self.backend, self.bucket, obj_key)) + _, _, obj_key = self._runtime_meta_refs(key) + logger.debug( + f"Uploading runtime metadata to: " + f"{self.backend}://{self.bucket}/{obj_key}" + ) self.storage.put_object(self.bucket, obj_key, json.dumps(runtime_meta)) if not is_lithops_worker(): - filename_local_path = os.path.join(CACHE_DIR, *path) + filename_local_path = self._local_runtime_meta_path(key) logger.debug( - "Storing runtime metadata into local cache: {}".format(filename_local_path)) - - if not os.path.exists(os.path.dirname(filename_local_path)): - os.makedirs(os.path.dirname(filename_local_path)) - - with open(filename_local_path, "w") as f: - f.write(json.dumps(runtime_meta)) + f"Storing runtime metadata into local cache: {filename_local_path}" + ) + self._write_runtime_meta_file(filename_local_path, runtime_meta) def delete_runtime_meta(self, key): - """ - Put the metadata given a runtime config. - :param runtime: name of the runtime - :param runtime_meta metadata - """ - path = [RUNTIMES_PREFIX, key + ".meta.json"] - obj_key = '/'.join(path).replace('\\', '/') - filename_local_path = os.path.join(CACHE_DIR, *path) + """Deletes the metadata of a runtime from storage and from the cache""" + _, _, obj_key = self._runtime_meta_refs(key) + filename_local_path = self._local_runtime_meta_path(key) if os.path.exists(filename_local_path): os.remove(filename_local_path) self.storage.delete_object(self.bucket, obj_key) diff --git a/lithops/storage/utils.py b/lithops/storage/utils.py index 23f1ed0a9..7c80bdc23 100644 --- a/lithops/storage/utils.py +++ b/lithops/storage/utils.py @@ -18,6 +18,8 @@ import os import time import logging +from typing import Any, List + from lithops.constants import JOBS_PREFIX @@ -33,19 +35,31 @@ class StorageNoSuchKeyError(Exception): - def __init__(self, bucket, key): + """Raised when a key a caller asked for is not in the storage backend""" + + def __init__(self, bucket: str, key: str): msg = f"No such key /{bucket}/{key} found in storage." - super(StorageNoSuchKeyError, self).__init__(msg) + super().__init__(msg) class StorageConfigMismatchError(Exception): - def __init__(self, current_path, prev_path): - msg = f"The data is stored at {prev_path}, but current storage is configured at {current_path}" - super(StorageConfigMismatchError, self).__init__(msg) + """ + Raised when the data of a previous run lives in a different backend or + bucket than the one currently configured + """ + + def __init__(self, current_path: List[str], prev_path: List[str]): + msg = ( + f"The data is stored at {prev_path}, but current storage " + f"is configured at {current_path}" + ) + super().__init__(msg) class CloudObject: - def __init__(self, backend, bucket, key): + """Reference to an object in a storage backend""" + + def __init__(self, backend: str, bucket: str, key: str): self.backend = backend self.bucket = bucket self.key = key @@ -56,7 +70,9 @@ def __str__(self): class CloudObjectUrl: - def __init__(self, url): + """Reference to an object named by its URL""" + + def __init__(self, url: str): self.url = url def __str__(self): @@ -64,7 +80,9 @@ def __str__(self): class CloudObjectLocal: - def __init__(self, path): + """Reference to an object that lives in the local filesystem""" + + def __init__(self, path: str): self.path = path self.bucket = os.path.dirname(path) self.key = os.path.basename(path) @@ -73,19 +91,24 @@ def __str__(self): return f'' -def clean_bucket(storage, bucket, prefix, sleep=5): +def clean_bucket( + storage: Any, bucket: str, prefix: str, sleep: int = 5 +) -> None: """ - Deletes all the files from COS. These files include the function, - the data serialization and the function invocation results. + Deletes every object under a prefix, which is where the serialized + function, its data and its results live. Lists again after each batch, + because a backend may report keys that the previous delete had not + applied yet """ msg = f"Deleting objects from bucket '{bucket}'" - msg = msg + f" and prefix '{prefix}'" if prefix else msg + if prefix: + msg = f"{msg} and prefix '{prefix}'" logger.info(msg) total_objects = 0 objects_to_delete = storage.list_keys(bucket, prefix) while objects_to_delete: - total_objects = total_objects + len(objects_to_delete) + total_objects += len(objects_to_delete) storage.delete_objects(bucket, objects_to_delete) time.sleep(sleep) objects_to_delete = storage.list_keys(bucket, prefix) @@ -93,85 +116,69 @@ def clean_bucket(storage, bucket, prefix, sleep=5): logger.info(f'Finished deleting objects, total found: {total_objects}') -def create_job_key(executor_id, job_id): - """ - Create job key - :param executor_id: prefix - :param job_id: Job's ID - :return: exec id - """ +def create_job_key(executor_id: str, job_id: str) -> str: + """Returns the key that identifies a job, shared by all of its calls""" return '-'.join([executor_id, job_id]) -def create_func_key(executor_id, function_hash): - """ - Create function key - :param prefix: prefix - :param executor_id: callset's ID - :return: function key - """ - return '/'.join([JOBS_PREFIX, executor_id, f'{function_hash}.{func_key_suffix}']) +def _jobs_key(*parts: str) -> str: + return '/'.join((JOBS_PREFIX, *parts)) -def create_data_key(executor_id, job_id): +def create_func_key(executor_id: str, function_hash: str) -> str: """ - Create aggregate data key - :param prefix: prefix - :param executor_id: callset's ID - :param job_id: Job's ID - :return: a key for aggregate data + Returns the key of a serialized function. The hash is part of the key, so + that the same function is uploaded only once per executor """ - job_key = create_job_key(executor_id, job_id) - return '/'.join([JOBS_PREFIX, job_key, agg_data_key_suffix]) + return _jobs_key(executor_id, f'{function_hash}.{func_key_suffix}') -def create_output_key(executor_id, job_id, call_id): - """ - Create output key - :param prefix: prefix - :param executor_id: Executor's ID - :param job_id: Job's ID - :param call_id: call's ID - :return: output key - """ - job_key = create_job_key(executor_id, job_id) - return '/'.join([JOBS_PREFIX, job_key, call_id, output_key_suffix]) +def create_data_key(executor_id: str, job_id: str) -> str: + """Returns the key of the aggregated data of every call of a job""" + return _jobs_key(create_job_key(executor_id, job_id), agg_data_key_suffix) -def create_status_key(executor_id, job_id, call_id): - """ - Create status key - :param prefix: prefix - :param executor_id: Executor's ID - :param job_id: Job's ID - :param call_id: call's ID - :return: status key - """ - job_key = create_job_key(executor_id, job_id) - return '/'.join([JOBS_PREFIX, job_key, call_id, status_key_suffix]) +def create_output_key(executor_id: str, job_id: str, call_id: str) -> str: + """Returns the key the result of a single call is written to""" + return _jobs_key( + create_job_key(executor_id, job_id), call_id, output_key_suffix + ) -def create_init_key(executor_id, job_id, call_id, act_id): +def create_status_key(executor_id: str, job_id: str, call_id: str) -> str: + """Returns the key the final status of a single call is written to""" + return _jobs_key( + create_job_key(executor_id, job_id), call_id, status_key_suffix + ) + + +def create_init_key( + executor_id: str, job_id: str, call_id: str, act_id: str +) -> str: """ - Create init key - :param prefix: prefix - :param executor_id: Executor's ID - :param job_id: Job's ID - :param call_id: call's ID - :return: output key + Returns the key a call writes when it starts running. The activation id + is part of the key, so that a retried call does not overwrite the mark of + the attempt that came before it """ - job_key = create_job_key(executor_id, job_id) - return '/'.join([JOBS_PREFIX, job_key, call_id, f'{act_id}{init_key_suffix}']) + return _jobs_key( + create_job_key(executor_id, job_id), + call_id, + f'{act_id}{init_key_suffix}', + ) -def get_storage_path(storage_config): +def get_storage_path(storage_config: dict) -> List[str]: + """Returns the backend and the bucket the data of a run lives in""" backend = storage_config['backend'] bucket = storage_config[backend]['storage_bucket'] - return [backend, bucket] -def check_storage_path(storage_config, prev_path): +def check_storage_path(storage_config: dict, prev_path: List[str]) -> None: + """ + Makes sure the configured storage is the one a previous run used, as data + written elsewhere is not reachable from here + """ current_path = get_storage_path(storage_config) if current_path != prev_path: raise StorageConfigMismatchError(current_path, prev_path) diff --git a/lithops/tests/conftest.py b/lithops/tests/conftest.py index 7fbcac6be..76d27f14a 100644 --- a/lithops/tests/conftest.py +++ b/lithops/tests/conftest.py @@ -18,6 +18,23 @@ def pytest_addoption(parser): parser.addoption("--region", metavar="", default=None, help="region") +@pytest.fixture(autouse=True) +def restore_environ(): + """ + Gives every test the environment back as it found it. + + Worker code sets process-wide variables of its own — LITHOPS_WORKER, the + session id, the monitoring queues — so a test that calls it leaks them + into every test that runs afterwards, and `monkeypatch.delenv` registers + no undo for a variable that was not there to begin with. That made the + outcome depend on the order the files happened to run in + """ + saved = os.environ.copy() + yield + os.environ.clear() + os.environ.update(saved) + + @pytest.fixture(scope="session", autouse=True) def setup_global(request): config = request.config diff --git a/lithops/tests/functions.py b/lithops/tests/functions.py index 7ac5ada40..12e4021e3 100644 --- a/lithops/tests/functions.py +++ b/lithops/tests/functions.py @@ -1,4 +1,5 @@ import lithops +import os import time import pickle @@ -193,3 +194,21 @@ def foo(self): def passthrough_function(x): return x.result + + +def raise_value_error(x): + raise ValueError('worker failed') + + +def sleep_seconds(x): + time.sleep(x) + return x + + +def echo_object(obj): + data = obj.data_stream.read() + return data.decode() if isinstance(data, bytes) else data + + +def echo_env_flag(x): + return os.environ.get('LITHOPS_TEST_FLAG') diff --git a/lithops/tests/test_config.py b/lithops/tests/test_config.py new file mode 100644 index 000000000..5d3a71252 --- /dev/null +++ b/lithops/tests/test_config.py @@ -0,0 +1,553 @@ +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import importlib +import json +from unittest.mock import MagicMock, patch + +import pytest +import yaml + +from lithops import constants as c +from lithops.config import ( + _ensure_lithops_section, + _resolve_mode_and_backend, + _section_with_user_agent, + default_config, + default_storage_config, + dump_yaml_config, + extract_localhost_config, + extract_serverless_config, + extract_standalone_config, + extract_storage_config, + get_default_config_filename, + get_log_info, + load_config, + load_yaml_config, +) +from lithops.version import __version__ + +_real_import_module = importlib.import_module + + +def _isolate_config_files(monkeypatch, tmp_path): + monkeypatch.delenv('LITHOPS_CONFIG', raising=False) + monkeypatch.delenv('LITHOPS_CONFIG_FILE', raising=False) + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(c, 'CONFIG_FILE', str(tmp_path / 'missing-user-config')) + monkeypatch.setattr(c, 'CONFIG_FILE_GLOBAL', str(tmp_path / 'missing-global-config')) + + +def _localhost_input(**lithops_extra): + lithops_cfg = {'mode': c.LOCALHOST, 'backend': c.LOCALHOST, 'storage': c.LOCALHOST} + lithops_cfg.update(lithops_extra) + return {'lithops': lithops_cfg} + + +class TestYamlConfig: + + def test_load_missing_file_returns_empty_dict(self, tmp_path): + assert load_yaml_config(str(tmp_path / 'does-not-exist.yml')) == {} + + def test_dump_and_load_roundtrip(self, tmp_path): + path = tmp_path / 'nested' / 'cfg.yml' + payload = {'lithops': {'mode': 'localhost'}, 'flag': True} + dump_yaml_config(str(path), payload) + loaded = load_yaml_config(str(path)) + assert loaded == payload + + def test_load_empty_file_returns_none(self, tmp_path): + path = tmp_path / 'empty.yml' + path.write_text('') + assert load_yaml_config(str(path)) is None + + def test_dump_filename_without_directory(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + dump_yaml_config('plain.yml', {'a': 1}) + assert load_yaml_config('plain.yml') == {'a': 1} + + +class TestConfigFileDiscovery: + + def test_env_lithops_config_file_wins(self, monkeypatch, tmp_path): + cfg = tmp_path / 'from-env.yml' + cfg.write_text('lithops: {}\n') + monkeypatch.setenv('LITHOPS_CONFIG_FILE', str(cfg)) + assert get_default_config_filename() == str(cfg) + + def test_env_lithops_config_file_returned_even_if_missing(self, monkeypatch, tmp_path): + missing = str(tmp_path / 'does-not-exist.yml') + monkeypatch.setenv('LITHOPS_CONFIG_FILE', missing) + assert get_default_config_filename() == missing + + def test_dotfile_in_cwd(self, monkeypatch, tmp_path): + _isolate_config_files(monkeypatch, tmp_path) + (tmp_path / '.lithops_config').write_text('lithops: {}\n') + assert get_default_config_filename() == str((tmp_path / '.lithops_config').resolve()) + + def test_user_config_file(self, monkeypatch, tmp_path): + _isolate_config_files(monkeypatch, tmp_path) + user_cfg = tmp_path / 'user-config' + user_cfg.write_text('lithops: {}\n') + monkeypatch.setattr(c, 'CONFIG_FILE', str(user_cfg)) + assert get_default_config_filename() == str(user_cfg) + + def test_global_config_file(self, monkeypatch, tmp_path): + _isolate_config_files(monkeypatch, tmp_path) + global_cfg = tmp_path / 'global-config' + global_cfg.write_text('lithops: {}\n') + monkeypatch.setattr(c, 'CONFIG_FILE_GLOBAL', str(global_cfg)) + assert get_default_config_filename() == str(global_cfg) + + def test_none_when_no_file_exists(self, monkeypatch, tmp_path): + _isolate_config_files(monkeypatch, tmp_path) + assert get_default_config_filename() is None + + +class TestLoadConfig: + + def test_explicit_missing_file_raises(self, tmp_path): + with pytest.raises(FileNotFoundError, match="doesn't exist"): + load_config(str(tmp_path / 'missing.yml')) + + def test_explicit_file(self, tmp_path): + cfg = tmp_path / 'cfg.yml' + cfg.write_text(yaml.dump({'lithops': {'mode': 'localhost'}})) + loaded = load_config(str(cfg), log=False) + assert loaded['lithops']['mode'] == 'localhost' + + def test_json_from_env(self, monkeypatch, tmp_path): + _isolate_config_files(monkeypatch, tmp_path) + monkeypatch.setenv('LITHOPS_CONFIG', json.dumps({'lithops': {'backend': 'localhost'}})) + loaded = load_config(log=False) + assert loaded['lithops']['backend'] == 'localhost' + + def test_fallback_to_localhost(self, monkeypatch, tmp_path): + _isolate_config_files(monkeypatch, tmp_path) + loaded = load_config(log=False) + assert loaded == { + 'lithops': { + 'mode': c.LOCALHOST, + 'backend': c.LOCALHOST, + 'storage': c.LOCALHOST, + } + } + + def test_fallback_does_not_share_mutable_state(self, monkeypatch, tmp_path): + _isolate_config_files(monkeypatch, tmp_path) + first = load_config(log=False) + first['lithops']['mode'] = 'serverless' + second = load_config(log=False) + assert second['lithops']['mode'] == c.LOCALHOST + + def test_json_env_wins_over_config_file(self, monkeypatch, tmp_path): + _isolate_config_files(monkeypatch, tmp_path) + (tmp_path / '.lithops_config').write_text(yaml.dump({'lithops': {'backend': 'fromfile'}})) + monkeypatch.setenv('LITHOPS_CONFIG', json.dumps({'lithops': {'backend': 'fromenv'}})) + loaded = load_config(log=False) + assert loaded['lithops']['backend'] == 'fromenv' + + def test_missing_env_config_file_falls_back_to_localhost(self, monkeypatch, tmp_path): + _isolate_config_files(monkeypatch, tmp_path) + monkeypatch.setenv('LITHOPS_CONFIG_FILE', str(tmp_path / 'nope.yml')) + loaded = load_config(log=False) + assert loaded['lithops'] == { + 'mode': c.LOCALHOST, 'backend': c.LOCALHOST, 'storage': c.LOCALHOST + } + + def test_explicit_empty_yaml_falls_back_to_localhost(self, tmp_path): + path = tmp_path / 'empty.yml' + path.write_text('') + loaded = load_config(str(path), log=False) + assert loaded['lithops']['backend'] == c.LOCALHOST + + def test_yaml_empty_mapping_falls_back_to_localhost(self, tmp_path): + path = tmp_path / 'empty-map.yml' + path.write_text('{}\n') + loaded = load_config(str(path), log=False) + assert loaded['lithops']['mode'] == c.LOCALHOST + + def test_file_with_empty_lithops_section_is_kept(self, tmp_path): + path = tmp_path / 'cfg.yml' + path.write_text('lithops: {}\n') + loaded = load_config(str(path), log=False) + assert loaded == {'lithops': {}} + + def test_json_env_empty_object_falls_back_to_localhost(self, monkeypatch, tmp_path): + _isolate_config_files(monkeypatch, tmp_path) + monkeypatch.setenv('LITHOPS_CONFIG', '{}') + loaded = load_config(log=False) + assert loaded['lithops']['mode'] == c.LOCALHOST + + def test_explicit_file_expands_user_home(self, monkeypatch, tmp_path): + cfg = tmp_path / 'home.yml' + cfg.write_text(yaml.dump({'lithops': {'mode': 'localhost'}})) + monkeypatch.setenv('HOME', str(tmp_path)) + loaded = load_config('~/home.yml', log=False) + assert loaded['lithops']['mode'] == 'localhost' + + +class TestGetLogInfo: + + def test_defaults(self): + level, fmt, stream, filename = get_log_info(config_data={'lithops': {}}) + assert level == c.LOGGER_LEVEL + assert fmt == c.LOGGER_FORMAT + assert stream == c.LOGGER_STREAM + assert filename is None + + def test_values_from_config(self): + level, fmt, stream, filename = get_log_info(config_data={ + 'lithops': { + 'log_level': 'DEBUG', + 'log_format': '%(message)s', + 'log_stream': 'ext://sys.stdout', + 'log_filename': '/tmp/lithops.log', + } + }) + assert level == 'DEBUG' + assert fmt == '%(message)s' + assert stream == 'ext://sys.stdout' + assert filename == '/tmp/lithops.log' + + def test_does_not_mutate_input(self): + src = {'lithops': {}} + get_log_info(config_data=src) + assert 'log_level' not in src['lithops'] + + def test_explicit_none_log_level_is_preserved(self): + level, *_ = get_log_info(config_data={'lithops': {'log_level': None}}) + assert level is None + + def test_missing_lithops_section_gets_defaults(self): + level, fmt, stream, filename = get_log_info(config_data={'other': 1}) + assert level == c.LOGGER_LEVEL + assert fmt == c.LOGGER_FORMAT + assert stream == c.LOGGER_STREAM + assert filename is None + + +class TestResolveModeAndBackend: + + def test_mode_without_backend_uses_mode_section(self): + cfg = { + 'lithops': {'mode': c.SERVERLESS}, + c.SERVERLESS: {'backend': 'code_engine'}, + } + backend, mode = _resolve_mode_and_backend(cfg) + assert mode == c.SERVERLESS + assert backend == 'code_engine' + + def test_mode_without_backend_uses_default_backend(self): + cfg = {'lithops': {'mode': c.LOCALHOST}} + backend, mode = _resolve_mode_and_backend(cfg) + assert backend == c.LOCALHOST + assert mode == c.LOCALHOST + + def test_backend_sets_mode(self): + cfg = {'lithops': {'backend': 'aws_lambda'}} + backend, mode = _resolve_mode_and_backend(cfg) + assert backend == 'aws_lambda' + assert mode == c.SERVERLESS + + def test_neither_uses_mode_default(self): + cfg = {'lithops': {}} + backend, mode = _resolve_mode_and_backend(cfg) + assert mode == c.MODE_DEFAULT + assert backend == c.SERVERLESS_BACKEND_DEFAULT + + def test_backend_wins_when_both_are_set(self): + cfg = {'lithops': {'mode': c.SERVERLESS, 'backend': c.LOCALHOST}} + backend, mode = _resolve_mode_and_backend(cfg) + assert backend == c.LOCALHOST + assert mode == c.LOCALHOST + + def test_standalone_backend_sets_standalone_mode(self): + cfg = {'lithops': {'backend': 'aws_ec2'}} + backend, mode = _resolve_mode_and_backend(cfg) + assert backend == 'aws_ec2' + assert mode == c.STANDALONE + + def test_mode_section_backend_none_is_still_applied(self): + cfg = { + 'lithops': {'mode': c.SERVERLESS}, + c.SERVERLESS: {'backend': None}, + } + backend, mode = _resolve_mode_and_backend(cfg) + assert backend is None + assert mode == c.SERVERLESS + + def test_empty_string_backend_is_treated_as_missing(self): + cfg = {'lithops': {'mode': c.LOCALHOST, 'backend': ''}} + backend, mode = _resolve_mode_and_backend(cfg) + assert backend == c.LOCALHOST + assert mode == c.LOCALHOST + + +class TestDefaultConfig: + + def test_localhost_completes_defaults(self): + cfg = default_config(config_data=_localhost_input()) + assert cfg['lithops']['mode'] == c.LOCALHOST + assert cfg['lithops']['backend'] == c.LOCALHOST + assert cfg['lithops']['storage'] == c.LOCALHOST + assert cfg['lithops']['chunksize'] == cfg['localhost']['worker_processes'] + assert cfg['lithops']['monitoring'] == 'storage' + assert cfg['lithops']['monitoring_interval'] == 0.1 + assert cfg['lithops']['execution_timeout'] == 3600 + assert 'localhost' in cfg + + def test_overwrite_lithops_and_backend_keys(self): + cfg = default_config( + config_data=_localhost_input(), + config_overwrite={ + 'lithops': {'monitoring_interval': 9}, + 'backend': {'runtime': 'python3'}, + }, + ) + assert cfg['lithops']['monitoring_interval'] == 9 + assert cfg['localhost']['runtime'] == 'python3' + + def test_does_not_mutate_input(self): + src = _localhost_input() + default_config(config_data=src) + assert 'chunksize' not in src['lithops'] + assert 'localhost' not in src + + def test_empty_lithops_section_is_replaced(self): + cfg = {'lithops': None} + _ensure_lithops_section(cfg) + assert cfg['lithops'] == {} + + def test_localhost_storage_rejected_for_other_backends(self): + fake_module = MagicMock() + fake_module.load_config.side_effect = lambda cfg: cfg.setdefault( + 'aws_lambda', {'worker_processes': 1} + ) + + def importer(name): + if name.endswith('aws_lambda.config'): + return fake_module + return _real_import_module(name) + + with patch('lithops.config.importlib.import_module', side_effect=importer): + with pytest.raises(Exception, match='Localhost storage backend cannot be used'): + default_config(config_data={ + 'lithops': {'backend': 'aws_lambda', 'storage': c.LOCALHOST}, + 'aws_lambda': {'worker_processes': 1}, + }) + + def test_standalone_sets_chunksize_zero(self): + fake_module = MagicMock() + + def load_config(cfg): + cfg.setdefault('standalone', {}) + cfg.setdefault('aws_ec2', {'worker_processes': 4}) + + fake_module.load_config.side_effect = load_config + + def importer(name): + if 'standalone.backends' in name: + return fake_module + return _real_import_module(name) + + with patch('lithops.config.importlib.import_module', side_effect=importer): + cfg = default_config( + config_data={ + 'lithops': {'backend': 'aws_ec2', 'storage': c.LOCALHOST}, + 'aws_ec2': {}, + }, + load_storage_config=False, + ) + + assert cfg['lithops']['mode'] == c.STANDALONE + assert cfg['lithops']['chunksize'] == 0 + fake_module.load_config.assert_called_once() + + def test_standalone_overwrites_user_chunksize(self): + fake_module = MagicMock() + fake_module.load_config.side_effect = lambda cfg: cfg.setdefault('standalone', {}) + + def importer(name): + if 'standalone.backends' in name: + return fake_module + return _real_import_module(name) + + with patch('lithops.config.importlib.import_module', side_effect=importer): + cfg = default_config( + config_data={ + 'lithops': {'backend': 'aws_ec2', 'chunksize': 8}, + 'aws_ec2': {'worker_processes': 4}, + }, + load_storage_config=False, + ) + assert cfg['lithops']['chunksize'] == 0 + + def test_user_chunksize_preserved_on_localhost(self): + cfg = default_config(config_data=_localhost_input(chunksize=8)) + assert cfg['lithops']['chunksize'] == 8 + + def test_backend_overwrite_worker_processes_sets_chunksize(self): + cfg = default_config( + config_data=_localhost_input(), + config_overwrite={'backend': {'worker_processes': 3}}, + ) + assert cfg['localhost']['worker_processes'] == 3 + assert cfg['lithops']['chunksize'] == 3 + + def test_empty_dict_config_data_loads_from_discovery(self, monkeypatch, tmp_path): + _isolate_config_files(monkeypatch, tmp_path) + cfg = default_config(config_data={}) + assert cfg['lithops']['mode'] == c.LOCALHOST + assert cfg['lithops']['backend'] == c.LOCALHOST + + def test_backend_overrides_conflicting_mode(self): + cfg = default_config(config_data={ + 'lithops': { + 'mode': c.SERVERLESS, + 'backend': c.LOCALHOST, + 'storage': c.LOCALHOST, + } + }) + assert cfg['lithops']['mode'] == c.LOCALHOST + assert cfg['lithops']['backend'] == c.LOCALHOST + + def test_overwrite_backend_rewrites_mode(self): + cfg = default_config( + config_data={'lithops': {'mode': c.SERVERLESS, 'storage': c.LOCALHOST}}, + config_overwrite={'lithops': {'backend': c.LOCALHOST}}, + ) + assert cfg['lithops']['mode'] == c.LOCALHOST + assert cfg['lithops']['backend'] == c.LOCALHOST + + def test_none_backend_section_is_replaced(self): + src = _localhost_input() + src['localhost'] = None + cfg = default_config(config_data=src) + assert isinstance(cfg['localhost'], dict) + assert cfg['localhost']['max_workers'] == 1 + + def test_empty_backend_overwrite_is_ignored(self): + cfg = default_config( + config_data=_localhost_input(), + config_overwrite={'backend': {}}, + ) + assert cfg['lithops']['backend'] == c.LOCALHOST + + def test_skip_storage_config_skips_localhost_storage_defaults(self): + cfg = default_config( + config_data=_localhost_input(), + load_storage_config=False, + ) + assert cfg['lithops']['monitoring_interval'] == 2 + assert 'storage_bucket' not in cfg.get('localhost', {}) + + def test_unknown_backend_raises(self): + with pytest.raises(Exception, match='Unknown compute backend'): + default_config(config_data={ + 'lithops': {'backend': 'not_a_backend', 'storage': c.LOCALHOST} + }) + + def test_unknown_mode_raises(self): + with pytest.raises(Exception, match='Unknown execution mode'): + default_config(config_data={ + 'lithops': {'mode': 'spaceship', 'storage': c.LOCALHOST} + }) + + +class TestStorageAndExtract: + + def test_default_storage_config_localhost(self): + cfg = default_storage_config(config_data=_localhost_input()) + assert cfg['lithops']['storage'] == c.LOCALHOST + assert cfg['localhost']['storage_bucket'] == 'storage' + + def test_default_storage_config_backend_override(self): + cfg = default_storage_config( + config_data={'lithops': {'storage': 'aws_s3'}}, + backend=c.LOCALHOST, + ) + assert cfg['lithops']['storage'] == c.LOCALHOST + + def test_extract_storage_config_sets_user_agent(self): + cfg = { + 'lithops': {'storage': c.LOCALHOST, 'monitoring_interval': 0.5}, + c.LOCALHOST: {'storage_bucket': 'storage'}, + } + extracted = extract_storage_config(cfg) + assert extracted['backend'] == c.LOCALHOST + assert extracted['monitoring_interval'] == 0.5 + assert extracted[c.LOCALHOST]['user_agent'] == f'lithops/{__version__}' + assert cfg[c.LOCALHOST]['user_agent'] == extracted[c.LOCALHOST]['user_agent'] + + def test_extract_storage_config_missing_backend_section(self): + cfg = {'lithops': {'storage': c.LOCALHOST}} + extracted = extract_storage_config(cfg) + assert extracted[c.LOCALHOST]['user_agent'] == f'lithops/{__version__}' + assert c.LOCALHOST not in cfg + + def test_extract_localhost_config_is_a_copy(self): + cfg = {c.LOCALHOST: {'runtime': 'python3', 'version': 2}} + extracted = extract_localhost_config(cfg) + extracted['runtime'] = 'other' + assert cfg[c.LOCALHOST]['runtime'] == 'python3' + + def test_extract_serverless_config(self): + cfg = { + 'lithops': {'backend': 'aws_lambda'}, + 'aws_lambda': {'region': 'us-east-1'}, + } + extracted = extract_serverless_config(cfg) + assert extracted['backend'] == 'aws_lambda' + assert extracted['aws_lambda']['region'] == 'us-east-1' + assert extracted['aws_lambda']['user_agent'] == f'lithops/{__version__}' + + def test_extract_standalone_config(self): + cfg = { + 'lithops': {'backend': 'aws_ec2', 'storage': c.LOCALHOST}, + c.STANDALONE: {'exec_mode': 'reuse'}, + 'aws_ec2': {'region': 'us-east-1'}, + } + extracted = extract_standalone_config(cfg) + assert extracted['backend'] == 'aws_ec2' + assert extracted['storage'] == c.LOCALHOST + assert extracted['exec_mode'] == 'reuse' + assert extracted['aws_ec2']['user_agent'] == f'lithops/{__version__}' + + def test_section_with_user_agent_uses_empty_dict_when_missing(self): + section = _section_with_user_agent({'lithops': {}}, 'aws_lambda') + assert section == {'user_agent': f'lithops/{__version__}'} + + def test_extract_does_not_mutate_empty_backend_section(self): + cfg = {'lithops': {'storage': c.LOCALHOST}, c.LOCALHOST: {}} + extract_storage_config(cfg) + assert 'user_agent' not in cfg[c.LOCALHOST] + + def test_extract_storage_config_default_monitoring_interval(self): + cfg = { + 'lithops': {'storage': c.LOCALHOST}, + c.LOCALHOST: {'storage_bucket': 'storage'}, + } + extracted = extract_storage_config(cfg) + assert extracted['monitoring_interval'] == c.LITHOPS_DEFAULT_CONFIG_KEYS['monitoring_interval'] + + def test_extract_standalone_does_not_share_standalone_section(self): + cfg = { + 'lithops': {'backend': 'aws_ec2', 'storage': c.LOCALHOST}, + c.STANDALONE: {'exec_mode': 'reuse'}, + 'aws_ec2': {'region': 'us-east-1'}, + } + extracted = extract_standalone_config(cfg) + extracted['exec_mode'] = 'consume' + assert cfg[c.STANDALONE]['exec_mode'] == 'reuse' diff --git a/lithops/tests/test_constants.py b/lithops/tests/test_constants.py new file mode 100644 index 000000000..4e7e9b342 --- /dev/null +++ b/lithops/tests/test_constants.py @@ -0,0 +1,117 @@ +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import os + +from lithops import constants +from lithops.utils import get_default_backend, get_mode + + +class TestConstants: + + def test_execution_modes(self): + assert constants.LOCALHOST == 'localhost' + assert constants.SERVERLESS == 'serverless' + assert constants.STANDALONE == 'standalone' + assert constants.MODE_DEFAULT == constants.SERVERLESS + + def test_default_backends_are_known(self): + assert constants.SERVERLESS_BACKEND_DEFAULT in constants.SERVERLESS_BACKENDS + assert constants.STANDALONE_BACKEND_DEFAULT in constants.STANDALONE_BACKENDS + assert constants.LOCALHOST not in constants.SERVERLESS_BACKENDS + assert constants.LOCALHOST not in constants.STANDALONE_BACKENDS + + def test_backend_collections_are_immutable(self): + assert isinstance(constants.SERVERLESS_BACKENDS, tuple) + assert isinstance(constants.STANDALONE_BACKENDS, tuple) + assert isinstance(constants.LOGGER_LEVEL_CHOICES, tuple) + + def test_temp_paths_are_under_lithops_temp_dir(self): + for path in ( + constants.JOBS_DIR, + constants.LOGS_DIR, + constants.MODULES_DIR, + constants.CUSTOM_RUNTIME_DIR, + constants.CLEANER_DIR, + constants.RN_LOG_FILE, + constants.SV_LOG_FILE, + constants.FN_LOG_FILE, + constants.SA_MASTER_LOG_FILE, + constants.SA_WORKER_LOG_FILE, + ): + assert path.startswith(constants.LITHOPS_TEMP_DIR) + assert os.path.isabs(path) + + def test_cleaner_files_are_under_cleaner_dir(self): + assert constants.CLEANER_PID_FILE.startswith(constants.CLEANER_DIR) + assert constants.CLEANER_LOG_FILE.startswith(constants.CLEANER_DIR) + + def test_config_paths(self): + assert constants.CONFIG_FILE.endswith(os.path.join('.lithops', 'config')) + assert constants.CACHE_DIR.endswith(os.path.join('.lithops', 'cache')) + assert constants.CONFIG_FILE_GLOBAL == '/etc/lithops/config' + + def test_local_temp_paths_use_native_separators(self): + assert constants.LITHOPS_TEMP_DIR == os.path.join( + constants.TEMP_DIR, constants.USER_TEMP_DIR + ) + assert constants.JOBS_DIR == os.path.join(constants.LITHOPS_TEMP_DIR, 'jobs') + assert constants.LOGS_DIR == os.path.join(constants.LITHOPS_TEMP_DIR, 'logs') + + def test_standalone_remote_paths_are_posix(self): + remote = ( + constants.SA_INSTALL_DIR, + constants.SA_SETUP_LOG_FILE, + constants.SA_SETUP_DONE_FILE, + constants.SA_CONFIG_FILE, + constants.SA_MASTER_DATA_FILE, + constants.SA_WORKER_DATA_FILE, + constants.CONFIG_FILE_GLOBAL, + ) + for path in remote: + assert path.startswith('/') + assert '\\' not in path + assert constants.SA_INSTALL_DIR == '/opt/lithops' + assert constants.SA_SETUP_LOG_FILE == '/opt/lithops/setup.log' + assert constants.SA_SETUP_DONE_FILE == '/opt/lithops/setup-done.flag' + assert constants.SA_CONFIG_FILE == '/opt/lithops/config' + assert constants.SA_MASTER_DATA_FILE == '/opt/lithops/master.data' + assert constants.SA_WORKER_DATA_FILE == '/opt/lithops/worker.data' + + def test_storage_prefixes_are_posix(self): + assert constants.JOBS_PREFIX == 'lithops.jobs' + assert constants.TEMP_PREFIX == 'lithops.jobs/tmp' + assert constants.LOGS_PREFIX == 'lithops.logs' + assert constants.RUNTIMES_PREFIX == 'lithops.runtimes' + assert '\\' not in constants.TEMP_PREFIX + + def test_default_config_keys(self): + assert set(constants.LITHOPS_DEFAULT_CONFIG_KEYS) == { + 'monitoring', 'monitoring_interval', 'execution_timeout' + } + assert constants.LITHOPS_DEFAULT_CONFIG_KEYS['monitoring_interval'] == 2 + + def test_get_mode_and_default_backend_round_trip(self): + assert get_mode(constants.LOCALHOST) == constants.LOCALHOST + assert get_mode(constants.SERVERLESS_BACKEND_DEFAULT) == constants.SERVERLESS + assert get_mode(constants.STANDALONE_BACKEND_DEFAULT) == constants.STANDALONE + assert get_default_backend(constants.LOCALHOST) == constants.LOCALHOST + assert get_default_backend(constants.SERVERLESS) == constants.SERVERLESS_BACKEND_DEFAULT + assert get_default_backend(constants.STANDALONE) == constants.STANDALONE_BACKEND_DEFAULT + + def test_every_known_backend_has_a_mode(self): + for backend in constants.SERVERLESS_BACKENDS: + assert get_mode(backend) == constants.SERVERLESS + for backend in constants.STANDALONE_BACKENDS: + assert get_mode(backend) == constants.STANDALONE diff --git a/lithops/tests/test_executors.py b/lithops/tests/test_executors.py new file mode 100644 index 000000000..4759ce75f --- /dev/null +++ b/lithops/tests/test_executors.py @@ -0,0 +1,596 @@ +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import os +import sys +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +import lithops +from lithops.constants import LOCALHOST, SERVERLESS, STANDALONE +from lithops.executors import ( + FunctionExecutor, + LocalhostExecutor, + ServerlessExecutor, + StandaloneExecutor, + _FixedModeExecutor, + _missing_plotting_extra, + _omit_none, +) +from lithops.storage.utils import create_job_key +from lithops.tests.conftest import TESTS_PREFIX +from lithops.tests.functions import ( + echo_object, + raise_value_error, + simple_map_function, + simple_reduce_function, + sleep_seconds, +) +from lithops.utils import FuturesList +from lithops.wait import ALL_COMPLETED, ALWAYS, ANY_COMPLETED + + +def _bare_executor(**attrs): + """Build a FunctionExecutor without running __init__ (no storage/backend).""" + executor = FunctionExecutor.__new__(FunctionExecutor) + defaults = { + 'config': {}, + 'futures': [], + 'cleaned_jobs': set(), + 'total_jobs': 0, + 'last_call': None, + 'data_cleaner': False, + 'executor_id': 'sess-0', + 'internal_storage': MagicMock(), + 'compute_handler': MagicMock(), + 'invoker': MagicMock(), + 'job_monitor': MagicMock(), + } + defaults.update(attrs) + for key, value in defaults.items(): + setattr(executor, key, value) + return executor + + +class FakeFuture: + def __init__(self, **kwargs): + self.done = False + self.success = False + self.error = False + self.futures = None + self._produce_output = True + self._read = False + self.job_id = 'M000' + self.job_key = 'sess-0/M000' + self.executor_id = 'sess-0' + self.function_name = 'fn' + self.runtime_memory = 256 + self.stats = {'worker_exec_time': 1.0} + self._result = None + self._exception_set = False + self._mapreduce = False + for key, value in kwargs.items(): + setattr(self, key, value) + + def result(self, throw_except=True, internal_storage=None): + return self._result + + def _set_mapreduce(self): + self._read = True + self._produce_output = False + self._mapreduce = True + + def _set_exception(self): + self._exception_set = True + + +class TestExecutorHelpers: + + def test_omit_none_keeps_falsey_non_none_values(self): + assert _omit_none({'a': 1, 'b': None, 'c': 0, 'd': False, 'e': ''}) == { + 'a': 1, 'c': 0, 'd': False, 'e': '' + } + + def test_missing_plotting_extra_mentions_method(self): + err = _missing_plotting_extra('plot') + assert isinstance(err, ModuleNotFoundError) + assert 'plot()' in str(err) + assert 'lithops[plotting]' in str(err) + + def test_build_config_overwrite_splits_lithops_and_backend(self): + overwrite = FunctionExecutor._build_config_overwrite( + mode='localhost', + backend=None, + storage='localhost', + monitoring=None, + kwargs={'runtime': 'python3', 'unused': None}, + ) + assert overwrite['lithops'] == {'mode': 'localhost', 'storage': 'localhost'} + assert overwrite['backend'] == {'runtime': 'python3'} + + def test_as_future_list_keeps_list_and_futures_list(self): + plain = [1, 2] + futures_list = FuturesList([1, 2]) + assert FunctionExecutor._as_future_list(plain) is plain + assert FunctionExecutor._as_future_list(futures_list) is futures_list + + def test_as_future_list_wraps_single_future(self): + future = FakeFuture() + assert FunctionExecutor._as_future_list(future) == [future] + + def test_disable_iterdata_output_only_for_futures_list(self): + future = FakeFuture(_produce_output=True) + FunctionExecutor._disable_iterdata_output([future]) + assert future._produce_output is True + + wrapped = FakeFuture(_produce_output=True) + FunctionExecutor._disable_iterdata_output(FuturesList([wrapped])) + assert wrapped._produce_output is False + + def test_create_job_id_increments_and_zero_fills(self): + executor = _bare_executor(total_jobs=0) + assert executor._create_job_id('A') == 'A000' + assert executor._create_job_id('M') == 'M001' + assert executor.total_jobs == 2 + + def test_create_compute_handler_localhost_versions(self): + with patch('lithops.executors.LocalhostHandlerV1') as v1, \ + patch('lithops.executors.extract_localhost_config', return_value={'version': 1}): + executor = _bare_executor(mode=LOCALHOST, config={'localhost': {}}) + assert executor._create_compute_handler() is v1.return_value + v1.assert_called_once() + + with patch('lithops.executors.LocalhostHandlerV2') as v2, \ + patch('lithops.executors.extract_localhost_config', return_value={}): + executor = _bare_executor(mode=LOCALHOST, config={'localhost': {}}) + assert executor._create_compute_handler() is v2.return_value + + def test_create_compute_handler_unknown_mode_returns_none(self): + executor = _bare_executor(mode='mystery', config={}) + assert executor._create_compute_handler() is None + + +class TestExecutorSubclasses: + + def test_fixed_mode_hierarchy(self): + assert issubclass(ServerlessExecutor, FunctionExecutor) + assert issubclass(StandaloneExecutor, FunctionExecutor) + assert issubclass(LocalhostExecutor, FunctionExecutor) + assert issubclass(ServerlessExecutor, _FixedModeExecutor) + assert ServerlessExecutor._mode == SERVERLESS + assert StandaloneExecutor._mode == STANDALONE + + @patch.object(FunctionExecutor, '__init__', return_value=None) + def test_serverless_executor_pins_mode(self, mock_init): + ServerlessExecutor(config={'lithops': {}}) + assert mock_init.call_args.kwargs['mode'] == SERVERLESS + + @patch.object(FunctionExecutor, '__init__', return_value=None) + def test_standalone_executor_pins_mode(self, mock_init): + StandaloneExecutor(config={'lithops': {}}) + assert mock_init.call_args.kwargs['mode'] == STANDALONE + + @patch.object(FunctionExecutor, '__init__', return_value=None) + def test_localhost_executor_pins_backend_and_storage(self, mock_init): + LocalhostExecutor() + assert mock_init.call_args.kwargs['backend'] == LOCALHOST + assert mock_init.call_args.kwargs['storage'] == LOCALHOST + + +class TestSubmitAndCleanup: + + def test_submit_map_uses_job_prefix_and_disables_futures_list_output(self): + executor = _bare_executor(total_jobs=0) + job = MagicMock() + submitted = [FakeFuture(_produce_output=True)] + iterdata = FuturesList(submitted) + + with patch.object(executor, '_run_map_job', return_value=(job, submitted)) as run_map: + job_id, out_job, out_fs = executor._submit_map( + lambda x: x, iterdata, job_prefix='A' + ) + + assert job_id == 'A000' + assert out_job is job + assert out_fs is submitted + assert submitted[0]._produce_output is False + assert run_map.call_args.kwargs['job_id'] == 'A000' + assert run_map.call_args.kwargs['iterdata'] is iterdata + + def test_cleanup_jobs_omits_exception_kwarg_on_success(self): + executor = _bare_executor() + future = FakeFuture() + with patch.object(executor, 'clean') as clean: + executor._cleanup_jobs([future]) + executor.compute_handler.clear.assert_called_once_with({future.job_key}) + clean.assert_called_once_with(clean_cloudobjects=False, force=False) + + def test_cleanup_jobs_passes_exception_and_force(self): + executor = _bare_executor() + future = FakeFuture() + error = RuntimeError('boom') + with patch.object(executor, 'clean') as clean: + executor._cleanup_jobs([future], exception=error, force=True) + executor.compute_handler.clear.assert_called_once_with( + {future.job_key}, exception=error + ) + clean.assert_called_once_with(clean_cloudobjects=False, force=True) + + def test_clean_does_not_wrap_futures_list(self): + future = FakeFuture(executor_id='abc-0', job_id='M000', done=True) + futures = FuturesList([future]) + executor = _bare_executor(cleaned_jobs=set(), executor_id='abc-0') + + with patch('lithops.executors._dump_cleaner_data') as dump, \ + patch('lithops.executors.sp.Popen'): + executor.clean(fs=futures, clean_cloudobjects=False) + + assert create_job_key('abc-0', 'M000') in executor.cleaned_jobs + dumped_jobs = dump.call_args_list[-1][0][0]['jobs_to_clean'] + assert create_job_key('abc-0', 'M000') in dumped_jobs + + def test_clean_fn_invalidates_function_cache(self): + from lithops.job.job import FUNCTION_CACHE + from lithops.storage.utils import create_func_key + + drop = create_func_key('abc-0', 'deadbeef') + keep = create_func_key('other-1', 'deadbeef') + saved = set(FUNCTION_CACHE) + FUNCTION_CACHE.update({drop, keep}) + try: + executor = _bare_executor(cleaned_jobs=set(), executor_id='abc-0') + with patch('lithops.executors._dump_cleaner_data'), \ + patch('lithops.executors.sp.Popen'): + executor.clean(clean_fn=True, clean_cloudobjects=False) + assert drop not in FUNCTION_CACHE + assert keep in FUNCTION_CACHE + finally: + FUNCTION_CACHE.clear() + FUNCTION_CACHE.update(saved) + + def test_dump_cleaner_data_recreates_missing_dir(self, tmp_path, monkeypatch): + import pickle + from lithops.executors import _dump_cleaner_data + + cleaner_dir = tmp_path / 'cleaner' + monkeypatch.setattr('lithops.executors.CLEANER_DIR', str(cleaner_dir)) + + _dump_cleaner_data({'jobs_to_clean': {'job-1'}}) + + dumped = list(cleaner_dir.iterdir()) + assert len(dumped) == 1 + with dumped[0].open('rb') as fh: + assert pickle.load(fh) == {'jobs_to_clean': {'job-1'}} + + +class TestWaitAndGetResult: + + @patch('lithops.executors.wait') + def test_wait_partitions_by_done_when_downloading_results(self, mock_wait): + finished = FakeFuture(done=True, success=True) + pending = FakeFuture(done=False, success=False) + executor = _bare_executor() + + done, notdone = executor.wait( + [finished, pending], download_results=True, show_progressbar=False + ) + + assert list(done) == [finished] + assert list(notdone) == [pending] + + @patch('lithops.executors.wait') + def test_wait_treats_success_as_done_when_not_downloading(self, mock_wait): + success = FakeFuture(done=False, success=True) + pending = FakeFuture(done=False, success=False) + executor = _bare_executor() + + done, notdone = executor.wait( + [success, pending], download_results=False, show_progressbar=False + ) + + assert list(done) == [success] + assert list(notdone) == [pending] + + @patch('lithops.executors.wait') + def test_wait_cleans_when_all_completed(self, mock_wait): + future = FakeFuture(done=True, success=True) + executor = _bare_executor(data_cleaner=True) + + with patch.object(executor, '_cleanup_jobs') as cleanup: + executor.wait([future], return_when=ALL_COMPLETED, show_progressbar=False) + + cleanup.assert_called_once() + assert cleanup.call_args.kwargs.get('force', False) is False + assert cleanup.call_args.kwargs.get('exception') is None + + @patch('lithops.executors.wait') + def test_wait_stops_monitor_when_all_tracked_futures_are_done(self, mock_wait): + future = FakeFuture(done=True, success=True) + executor = _bare_executor(futures=[future]) + executor.wait([future], return_when=ALL_COMPLETED, show_progressbar=False) + executor.job_monitor.stop.assert_called_once() + + @patch('lithops.executors.wait') + def test_wait_keeps_monitor_when_other_futures_are_pending(self, mock_wait): + done = FakeFuture(done=True, success=True) + pending = FakeFuture(done=False, success=False) + executor = _bare_executor(futures=[done, pending]) + executor.wait([done], return_when=ALL_COMPLETED, show_progressbar=False) + executor.job_monitor.stop.assert_not_called() + + @patch('lithops.executors.wait', side_effect=RuntimeError('boom')) + def test_wait_exception_stops_invoker_and_reraises(self, mock_wait): + future = FakeFuture() + executor = _bare_executor(data_cleaner=True) + + with patch.object(executor, '_cleanup_jobs') as cleanup: + with pytest.raises(RuntimeError, match='boom'): + executor.wait([future], show_progressbar=False) + + executor.invoker.stop.assert_called_once() + executor.job_monitor.remove.assert_called_once() + assert future._exception_set is True + assert cleanup.call_args.kwargs['force'] is True + assert isinstance(cleanup.call_args.kwargs['exception'], RuntimeError) + + def test_get_result_unwraps_single_non_map_result(self): + future = FakeFuture(_result=42) + executor = _bare_executor(last_call='call_async', futures=[future]) + + with patch.object(executor, 'wait', return_value=([future], [])): + assert executor.get_result() == 42 + assert future._read is True + + def test_get_result_keeps_list_for_map(self): + future = FakeFuture(_result=42) + executor = _bare_executor(last_call='map', futures=[future]) + + with patch.object(executor, 'wait', return_value=([future], [])): + assert executor.get_result() == [42] + + def test_get_result_skips_nested_and_already_read_futures(self): + nested = FakeFuture(futures=[FakeFuture()], _result='nested') + consumed = FakeFuture(_read=True, _result='old') + pending = FakeFuture(_result='new') + executor = _bare_executor( + last_call='map_reduce', + futures=[nested, consumed, pending], + ) + + with patch.object( + executor, 'wait', return_value=([nested, consumed, pending], []) + ): + assert executor.get_result() == 'new' + + def test_plot_returns_when_no_ready_futures(self): + executor = _bare_executor(futures=[FakeFuture(success=False, done=False)]) + assert executor.plot() is None + + def test_plot_calls_timeline_and_histogram(self): + future = FakeFuture(success=True, done=True, error=False) + executor = _bare_executor(futures=[future]) + fake_plots = MagicMock() + with patch.dict(sys.modules, {'lithops.plots': fake_plots}): + executor.plot(dst='/tmp/out', figsize=(4, 3)) + fake_plots.create_timeline.assert_called_once_with( + [future], '/tmp/out', (4, 3) + ) + fake_plots.create_histogram.assert_called_once_with( + [future], '/tmp/out', (4, 3) + ) + + def test_plot_missing_extra_raises(self): + future = FakeFuture(success=True, done=True, error=False) + executor = _bare_executor(futures=[future]) + with patch.dict(sys.modules, {'lithops.plots': None}): + with pytest.raises(ModuleNotFoundError, match=r'plot\(\)'): + executor.plot() + + def test_job_summary_warns_when_backend_has_no_calc_cost(self): + pytest.importorskip('pandas') + executor = _bare_executor() + executor.compute_handler.backend = SimpleNamespace(name='localhost') + with patch('lithops.executors.logger.warning') as warn: + executor.job_summary() + warn.assert_called_once() + assert "isn't supported" in warn.call_args[0][0] + + def test_job_summary_writes_csv_when_backend_supports_cost( + self, tmp_path, monkeypatch + ): + pytest.importorskip('pandas') + monkeypatch.setattr('lithops.executors.constants.LOGS_DIR', str(tmp_path)) + backend = MagicMock() + backend.calc_cost.return_value = 1.25 + executor = _bare_executor( + futures=[ + FakeFuture( + job_id='M000', function_name='fn', runtime_memory=128, + stats={'worker_exec_time': 0.5}, + ), + FakeFuture( + job_id='M000', function_name='fn', runtime_memory=128, + stats={'worker_exec_time': 1.5}, + ), + ] + ) + executor.compute_handler.backend = backend + executor.log_path = None + executor.job_summary() + assert executor.log_path + assert os.path.exists(executor.log_path) + text = open(executor.log_path).read() + assert 'Summary' in text + assert 'M000' in text + + def test_map_reduce_always_skips_waiting_for_map(self): + executor = _bare_executor(total_jobs=0) + map_futures = [FakeFuture(job_id='M000')] + reduce_futures = [FakeFuture(job_id='R000')] + job = MagicMock() + with patch.object( + executor, '_submit_map', return_value=('M000', job, map_futures) + ), patch.object(executor, 'wait') as wait, patch.object( + executor, '_run_reduce_job', return_value=reduce_futures + ): + result = executor.map_reduce( + lambda x: x, [1], lambda xs: xs, spawn_reducer=ALWAYS + ) + wait.assert_not_called() + assert map_futures[0]._mapreduce is True + assert list(result) == map_futures + reduce_futures + + +class TestExecutorLocalhost: + """Live localhost checks for the refactored public API.""" + + def test_call_async_job_prefix_and_unwrapped_result(self): + fexec = lithops.FunctionExecutor(config=pytest.lithops_config) + future = fexec.call_async(lambda x: x + 1, 1) + assert future.job_id.startswith('A') + assert fexec.last_call == 'call_async' + assert fexec.get_result() == 2 + + def test_map_job_prefix_and_list_result(self): + fexec = lithops.FunctionExecutor(config=pytest.lithops_config) + futures = fexec.map(lambda x: x * 2, [1, 2, 3]) + assert all(future.job_id.startswith('M') for future in futures) + assert fexec.last_call == 'map' + assert fexec.get_result() == [2, 4, 6] + + def test_map_reduce_map_and_reduce_job_ids(self): + fexec = lithops.FunctionExecutor(config=pytest.lithops_config) + futures = fexec.map_reduce( + simple_map_function, + [(1, 1), (2, 2)], + simple_reduce_function, + ) + job_ids = {future.job_id for future in futures} + assert any(job_id.startswith('M') for job_id in job_ids) + assert any(job_id.startswith('R') for job_id in job_ids) + assert fexec.last_call == 'map_reduce' + assert fexec.get_result() == 6 + + def test_map_reduce_spawn_reducer_always_and_percentage(self): + fexec = lithops.FunctionExecutor(config=pytest.lithops_config) + futures = fexec.map_reduce( + simple_map_function, + [(1, 1), (2, 2)], + simple_reduce_function, + spawn_reducer=ALWAYS, + ) + assert any(future.job_id.startswith('R') for future in futures) + assert fexec.get_result() == 6 + + fexec = lithops.FunctionExecutor(config=pytest.lithops_config) + futures = fexec.map_reduce( + simple_map_function, + [(1, 1), (2, 2), (3, 3), (4, 4)], + simple_reduce_function, + spawn_reducer=50, + ) + assert fexec.get_result() == 20 + + def test_map_local_file_partitions(self, tmp_path): + path = tmp_path / 'data.txt' + text = 'alpha beta gamma delta\n' * 20 + path.write_text(text) + fexec = lithops.FunctionExecutor(config=pytest.lithops_config) + futures = fexec.map( + echo_object, str(path), obj_chunk_number=2, obj_newline=None + ) + result = fexec.get_result() + assert len(futures) == 2 + assert ''.join(result) == text + + def test_map_local_file_chunk_size(self, tmp_path): + path = tmp_path / 'data.txt' + text = 'x' * 100 + path.write_text(text) + fexec = lithops.FunctionExecutor(config=pytest.lithops_config) + futures = fexec.map( + echo_object, str(path), obj_chunk_size=40, obj_newline=None + ) + result = fexec.get_result() + assert len(futures) == 3 + assert ''.join(result) == text + + def test_map_local_directory(self, tmp_path): + folder = tmp_path / 'files' + folder.mkdir() + (folder / 'a.txt').write_text('one') + (folder / 'b.txt').write_text('two') + fexec = lithops.FunctionExecutor(config=pytest.lithops_config) + fexec.map(echo_object, str(folder)) + assert sorted(fexec.get_result()) == ['one', 'two'] + + def test_context_manager_and_localhost_executor(self): + with lithops.LocalhostExecutor(config=pytest.lithops_config) as fexec: + fexec.map(simple_map_function, [(3, 4)]) + assert fexec.get_result() == [7] + assert fexec.mode == LOCALHOST + assert fexec.backend == LOCALHOST + + def test_get_result_reraises_worker_exception(self): + fexec = lithops.FunctionExecutor(config=pytest.lithops_config) + fexec.map(raise_value_error, [1]) + with pytest.raises(ValueError, match='worker failed'): + fexec.get_result() + + def test_get_result_throw_except_false_does_not_reraise(self): + fexec = lithops.FunctionExecutor(config=pytest.lithops_config) + fexec.map(raise_value_error, [1]) + fexec.get_result(throw_except=False) + + def test_wait_any_completed_and_percentage(self): + fexec = lithops.FunctionExecutor(config=pytest.lithops_config) + futures = fexec.map(sleep_seconds, [0, 2, 2]) + done, notdone = fexec.wait(futures, return_when=ANY_COMPLETED) + assert len(done) >= 1 + fexec.wait(futures) + assert fexec.get_result() == [0, 2, 2] + + fexec = lithops.FunctionExecutor(config=pytest.lithops_config) + futures = fexec.map(sleep_seconds, [0, 0, 3, 3]) + done, notdone = fexec.wait(futures, return_when=50) + assert len(done) >= 2 + fexec.wait() + + def test_execution_timeout(self): + fexec = lithops.FunctionExecutor(config=pytest.lithops_config) + fexec.map(sleep_seconds, [30], timeout=3) + with pytest.raises(Exception): + fexec.get_result() + + def test_clean_after_job(self): + fexec = lithops.FunctionExecutor(config=pytest.lithops_config) + fexec.map(lambda x: x, [1, 2]) + assert fexec.get_result() == [1, 2] + fexec.clean(force=True) + assert fexec.cleaned_jobs + + def test_map_obj_parameter_over_storage_prefix(self): + fexec = lithops.FunctionExecutor(config=pytest.lithops_config) + bucket = fexec.storage.bucket + prefix = TESTS_PREFIX + '/echo-obj/' + key = prefix + 'a.txt' + fexec.storage.put_object(bucket, key, b'alpha') + url = fexec.config['lithops']['storage'] + '://' + bucket + '/' + prefix + try: + fexec.map(echo_object, url) + assert fexec.get_result() == ['alpha'] + finally: + fexec.storage.delete_object(bucket, key) diff --git a/lithops/tests/test_future.py b/lithops/tests/test_future.py index 4d18dd9f2..9f245eaef 100644 --- a/lithops/tests/test_future.py +++ b/lithops/tests/test_future.py @@ -1,6 +1,27 @@ +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import base64 +import pickle +import zlib +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + import pytest import lithops +from lithops.future import ResponseFuture, _pickle_from_encoded, _stats_from_prefixed_keys class HasAmbiguousTruthValue: @@ -16,6 +37,377 @@ def __bool__(self): ) +STORAGE_CONFIG = { + 'backend': 'localhost', + 'localhost': {'storage_bucket': 'test-bucket'}, +} + + +def _job(**overrides): + values = dict( + job_id='M000', + job_key='sess-0/M000', + executor_id='sess-0', + function_name='fn', + execution_timeout=300, + runtime_name='python', + runtime_memory=256, + ) + values.update(overrides) + return SimpleNamespace(**values) + + +def _future(job_metadata=None, **job_kwargs): + return ResponseFuture('00000', _job(**job_kwargs), job_metadata or {}, STORAGE_CONFIG) + + +class TestStatsHelper: + + def test_keeps_func_host_and_worker_prefixes(self): + mapping = { + 'func_name': 'fn', + 'host_submit_tstamp': 1.0, + 'worker_start_tstamp': 2.0, + 'activation_id': 'skip-me', + 'type': '__end__', + } + assert _stats_from_prefixed_keys(mapping) == { + 'func_name': 'fn', + 'host_submit_tstamp': 1.0, + 'worker_start_tstamp': 2.0, + } + + def test_empty_mapping(self): + assert _stats_from_prefixed_keys({}) == {} + + +class TestResponseFutureState: + + def test_new_future_collects_job_metadata_stats(self): + future = _future({'func_name': 'fn', 'ignored': True}) + assert future.new + assert not future.invoked + assert not future.ready + assert not future.success + assert not future.done + assert not future.error + assert future.stats['func_name'] == 'fn' + assert 'ignored' not in future.stats + + def test_success_includes_error_state(self): + future = _future() + future._set_state(ResponseFuture.State.Error) + assert future.error + assert future.success + assert future.done + + def test_done_includes_unknown(self): + future = _future() + future._set_state(ResponseFuture.State.Unknown) + assert future.done + assert not future.success + + def test_set_running_and_ready(self): + future = _future() + future._set_running({'activation_id': 'act-1'}) + assert future.running + assert future.activation_id == 'act-1' + future._set_ready({'activation_id': 'act-1', 'type': '__end__'}) + assert future.ready + + def test_set_mapreduce_marks_successful_future_done(self): + future = _future() + future._set_state(ResponseFuture.State.Success) + future._set_mapreduce() + assert future.done + assert future._produce_output is False + assert future._read is True + + def test_set_mapreduce_does_not_advance_unsuccessful_future(self): + future = _future() + future._set_invoked() + future._set_mapreduce() + assert future.invoked + assert future._produce_output is False + + def test_status_and_result_reject_new_state(self): + future = _future() + with pytest.raises(ValueError, match='task not yet invoked'): + future.status() + with pytest.raises(ValueError, match='Task not yet invoked'): + future.result() + + def test_cancel_is_not_implemented(self): + future = _future() + with pytest.raises(NotImplementedError): + future.cancel() + with pytest.raises(NotImplementedError): + future.cancelled() + + def test_status_returns_cached_call_status_when_done(self): + future = _future() + future._call_status = {'already': 'done'} + future._set_state(ResponseFuture.State.Done) + assert future.status() == {'already': 'done'} + + def test_write_activation_logs(self, tmp_path, monkeypatch): + monkeypatch.setattr('lithops.future.LOGS_DIR', str(tmp_path)) + fn_log = tmp_path / 'functions.log' + monkeypatch.setattr('lithops.future.FN_LOG_FILE', str(fn_log)) + + future = _future() + future.activation_id = 'act-9' + raw = 'hello\nworld\n' + future._call_status = { + 'logs': base64.b64encode(zlib.compress(raw.encode())).decode(), + } + + future._write_activation_logs() + + job_log = (tmp_path / 'sess-0-M000.log').read_text() + assert "Activation: 'python' (act-9)" in job_log + assert 'hello' in job_log + assert fn_log.read_text() == job_log + + def test_write_activation_logs_recreates_missing_log_dir(self, tmp_path, monkeypatch): + logs_dir = tmp_path / 'missing' / 'logs' + monkeypatch.setattr('lithops.future.LOGS_DIR', str(logs_dir)) + fn_log = tmp_path / 'missing' / 'functions.log' + monkeypatch.setattr('lithops.future.FN_LOG_FILE', str(fn_log)) + + future = _future() + future.activation_id = 'act-9' + raw = 'hello\nworld\n' + future._call_status = { + 'logs': base64.b64encode(zlib.compress(raw.encode())).decode(), + } + + future._write_activation_logs() + + assert (logs_dir / 'sess-0-M000.log').exists() + assert fn_log.exists() + + def test_set_exception_marks_unfinished_future_unknown(self): + future = _future() + future._set_invoked() + future._set_exception() + assert future._read is True + assert future.done + assert future._state == ResponseFuture.State.Unknown + + def test_set_exception_does_not_change_already_done_state(self): + future = _future() + future._set_state(ResponseFuture.State.Done) + future._set_exception() + assert future._state == ResponseFuture.State.Done + assert future._read is True + + def test_futures_property_tracks_new_futures(self): + future = _future() + assert future.futures is False + future._new_futures = [] + assert future.futures is True + + +def _encode(obj): + return str(pickle.dumps(obj)) + + +def _end_status(**overrides): + status = { + 'type': '__end__', + 'exception': False, + 'activation_id': 'act-1', + 'func_result_size': 1, + 'worker_start_tstamp': 1.0, + 'worker_end_tstamp': 2.5, + 'host_submit_tstamp': 0.5, + } + status.update(overrides) + return status + + +class TestPickleHelper: + + def test_roundtrip(self): + assert _pickle_from_encoded(_encode({'k': 1})) == {'k': 1} + + +class TestResponseFutureStatusAndResult: + + def test_poll_check_only_returns_without_waiting(self): + future = _future() + future._set_invoked() + storage = MagicMock() + storage.get_storage_config.return_value = STORAGE_CONFIG + storage.get_call_status.return_value = None + assert future.status(internal_storage=storage, check_only=True) is None + storage.get_call_status.assert_called_once() + + def test_poll_retries_until_status_appears(self, monkeypatch): + future = _future() + future._set_invoked() + storage = MagicMock() + storage.get_storage_config.return_value = STORAGE_CONFIG + storage.get_call_status.side_effect = [None, None, _end_status(result=_encode(9))] + monkeypatch.setattr('lithops.future.time.sleep', lambda *_: None) + status = future.status(internal_storage=storage, wait_dur_sec=0) + assert storage.get_call_status.call_count == 3 + assert status['activation_id'] == 'act-1' + assert future.done + assert future.result(internal_storage=storage) == 9 + + def test_status_creates_internal_storage_when_missing(self): + future = _future() + future._set_invoked() + with patch('lithops.future.InternalStorage') as storage_cls: + inst = storage_cls.return_value + inst.get_storage_config.return_value = STORAGE_CONFIG + inst.get_call_status.return_value = _end_status(func_result_size=0) + future.status() + storage_cls.assert_called_once_with(STORAGE_CONFIG) + assert future._produce_output is False + assert future.done + + def test_status_refetches_init_call_status(self): + future = _future() + future._set_invoked() + future._call_status = {'type': '__init__', 'activation_id': 'boot'} + storage = MagicMock() + storage.get_storage_config.return_value = STORAGE_CONFIG + storage.get_call_status.return_value = _end_status(result=_encode('ok')) + future.status(internal_storage=storage) + assert future.activation_id == 'act-1' + assert future.done + + def test_exception_is_reraised_by_default(self): + future = _future() + future._set_invoked() + storage = MagicMock() + storage.get_storage_config.return_value = STORAGE_CONFIG + storage.get_call_status.return_value = _end_status( + exception=True, + exc_info=_encode((ValueError, ValueError('boom'), None)), + ) + with pytest.raises(ValueError, match='boom'): + future.status(internal_storage=storage) + assert future.error + + def test_exception_can_be_suppressed(self): + future = _future() + future._set_invoked() + storage = MagicMock() + storage.get_storage_config.return_value = STORAGE_CONFIG + storage.get_call_status.return_value = _end_status( + exception=True, + exc_info=_encode((ValueError, ValueError('boom'), None)), + ) + assert future.status(internal_storage=storage, throw_except=False) is None + assert future.error + + def test_handler_exception_strips_marker_argument(self): + future = _future() + future._set_invoked() + storage = MagicMock() + storage.get_storage_config.return_value = STORAGE_CONFIG + storage.get_call_status.return_value = _end_status( + exception=True, + exc_info=_encode((Exception, Exception('HANDLER', 'inner'), None)), + ) + with pytest.raises(Exception, match='inner'): + future.status(internal_storage=storage) + assert future._handler_exception is True + + def test_pickle_fail_wraps_exception_dict(self): + future = _future() + future._set_invoked() + storage = MagicMock() + storage.get_storage_config.return_value = STORAGE_CONFIG + storage.get_call_status.return_value = _end_status( + exception=True, + exc_pickle_fail=True, + exc_info=_encode({'exc_value': 'pickle-broke', 'exc_traceback': None}), + ) + with pytest.raises(Exception, match='pickle-broke'): + future.status(internal_storage=storage) + + def test_new_futures_wraps_a_single_response_future(self): + nested = _future() + future = _future() + future._set_invoked() + storage = MagicMock() + storage.get_storage_config.return_value = STORAGE_CONFIG + storage.get_call_status.return_value = _end_status( + new_futures=_encode(nested), + func_result_size=0, + ) + future.status(internal_storage=storage) + assert len(future._new_futures) == 1 + assert isinstance(future._new_futures[0], ResponseFuture) + assert future._new_futures[0].call_id == nested.call_id + assert future.result(internal_storage=storage) == future._new_futures + + def test_new_futures_keeps_a_list(self): + nested = [_future(), _future()] + future = _future() + future._set_invoked() + storage = MagicMock() + storage.get_storage_config.return_value = STORAGE_CONFIG + storage.get_call_status.return_value = _end_status( + new_futures=_encode(nested), + ) + future.status(internal_storage=storage) + assert isinstance(future._new_futures, list) + assert len(future._new_futures) == 2 + assert {f.call_id for f in future._new_futures} == {nested[0].call_id, nested[1].call_id} + + def test_result_polls_output_then_unpickles(self, monkeypatch): + future = _future() + future._set_invoked() + future._call_status = _end_status() + future._set_state(ResponseFuture.State.Success) + storage = MagicMock() + storage.get_call_output.side_effect = [None, pickle.dumps('later')] + monkeypatch.setattr('lithops.future.time.sleep', lambda *_: None) + assert future.result(internal_storage=storage, retries=5, wait_dur_sec=0) == 'later' + assert future.done + assert future.stats['host_result_query_count'] == 2 + + def test_result_missing_output_raises_by_default(self, monkeypatch): + future = _future() + future._set_invoked() + future._call_status = _end_status() + future._set_state(ResponseFuture.State.Success) + storage = MagicMock() + storage.get_call_output.return_value = None + monkeypatch.setattr('lithops.future.time.sleep', lambda *_: None) + with pytest.raises(Exception, match='Unable to get the result'): + future.result(internal_storage=storage, retries=2, wait_dur_sec=0) + + def test_result_missing_output_can_be_suppressed(self, monkeypatch): + future = _future() + future._set_invoked() + future._call_status = _end_status() + future._set_state(ResponseFuture.State.Success) + storage = MagicMock() + storage.get_call_output.return_value = None + monkeypatch.setattr('lithops.future.time.sleep', lambda *_: None) + assert future.result( + internal_storage=storage, retries=1, wait_dur_sec=0, throw_except=False + ) is None + assert future.error + + def test_result_creates_storage_when_not_done(self): + future = _future() + future._set_invoked() + with patch('lithops.future.InternalStorage') as storage_cls: + inst = storage_cls.return_value + inst.get_storage_config.return_value = STORAGE_CONFIG + inst.get_call_status.return_value = _end_status(result=_encode('x')) + assert future.result() == 'x' + storage_cls.assert_called_once() + + def test_fn_returns_obj_with_ambiguous_truth_value(): def returns_obj_with_ambiguous_truth_value(param): return HasAmbiguousTruthValue(param) diff --git a/lithops/tests/test_invokers.py b/lithops/tests/test_invokers.py new file mode 100644 index 000000000..41e925d53 --- /dev/null +++ b/lithops/tests/test_invokers.py @@ -0,0 +1,972 @@ +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import os +import queue +import sys +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from lithops.constants import LOCALHOST, SERVERLESS +from lithops.future import ResponseFuture +from lithops.invokers import ( + BatchInvoker, + FaaSInvoker, + Invoker, + _format_call_id, + _timed_invoke, + _verify_runtime_meta, + create_invoker, + extend_runtime, +) +from lithops.utils import BackendType, version_str +from lithops.version import __version__ + + +class TestInvokerHelpers: + + def test_format_call_id_zero_fills(self): + assert _format_call_id(0) == '00000' + assert _format_call_id(12) == '00012' + assert _format_call_id(99999) == '99999' + + def test_verify_runtime_meta_accepts_matching_versions(self): + _verify_runtime_meta( + { + 'lithops_version': __version__, + 'python_version': version_str(sys.version_info), + }, + 'python:3', + ) + + def test_verify_runtime_meta_lithops_mismatch(self): + with pytest.raises(Exception, match='Lithops version mismatch'): + _verify_runtime_meta( + { + 'lithops_version': '0.0.0', + 'python_version': version_str(sys.version_info), + }, + 'python:3', + ) + + def test_verify_runtime_meta_python_mismatch_includes_runtime_name(self): + with pytest.raises(Exception, match="indicated runtime 'my-runtime'"): + _verify_runtime_meta( + { + 'lithops_version': __version__, + 'python_version': '2.7', + }, + 'my-runtime', + ) + + +class TestCreateInvoker: + + @patch('lithops.invokers.BatchInvoker') + def test_creates_batch_invoker(self, batch_cls): + handler = MagicMock() + handler.get_backend_type.return_value = BackendType.BATCH.value + result = create_invoker('cfg', 'ex', 'store', handler, 'monitor') + batch_cls.assert_called_once_with('cfg', 'ex', 'store', handler, 'monitor') + assert result is batch_cls.return_value + + @patch('lithops.invokers.FaaSInvoker') + def test_creates_faas_invoker(self, faas_cls): + handler = MagicMock() + handler.get_backend_type.return_value = BackendType.FAAS.value + result = create_invoker('cfg', 'ex', 'store', handler, 'monitor') + faas_cls.assert_called_once_with('cfg', 'ex', 'store', handler, 'monitor') + assert result is faas_cls.return_value + + def test_unknown_backend_type_returns_none(self): + handler = MagicMock() + handler.get_backend_type.return_value = 'mystery' + assert create_invoker('cfg', 'ex', 'store', handler, 'monitor') is None + + +def _matching_runtime_meta(): + return { + 'lithops_version': __version__, + 'python_version': version_str(sys.version_info), + 'runtime_timeout': 300, + } + + +def _job(**overrides): + values = dict( + executor_id='sess-0', + job_id='M000', + job_key='sess-0/M000', + function_name='fn', + func_key='fk', + data_key='dk', + extra_env=None, + total_calls=2, + execution_timeout=60, + data_byte_ranges=[(0, 1), (2, 3)], + chunksize=1, + worker_processes=1, + runtime_name='python:3', + runtime_memory=256, + metadata={'func_name': 'fn'}, + ) + values.update(overrides) + return SimpleNamespace(**values) + + +def _bare_invoker(**attrs): + inv = Invoker.__new__(Invoker) + defaults = dict( + executor_id='sess-0', + mode=SERVERLESS, + runtime_name='python:3', + runtime_info={ + 'runtime_name': 'python:3', + 'runtime_memory': 256, + 'runtime_timeout': 300, + 'max_workers': 8, + }, + compute_handler=MagicMock(), + internal_storage=MagicMock(), + config={'lithops': {'mode': SERVERLESS, 'backend': 'ibm_cf'}, 'ibm_cf': {}}, + backend='ibm_cf', + include_function=False, + prometheus=MagicMock(), + job_monitor=MagicMock(), + storage_config={'backend': 'localhost', 'localhost': {'storage_bucket': 'test-bucket'}}, + max_workers=8, + log_level='INFO', + ) + defaults.update(attrs) + for key, value in defaults.items(): + setattr(inv, key, value) + return inv + + +class TestTimedInvoke: + + def test_timed_invoke_returns_activation_and_duration(self): + handler = MagicMock() + handler.invoke.return_value = 'act-9' + activation_id, resp_time = _timed_invoke(handler, {'x': 1}) + assert activation_id == 'act-9' + assert isinstance(resp_time, str) + handler.invoke.assert_called_once_with({'x': 1}) + + +class TestSelectRuntime: + + def test_serverless_uses_override_memory_and_skips_deploy_when_meta_exists(self): + inv = _bare_invoker() + inv.internal_storage.get_runtime_meta.return_value = _matching_runtime_meta() + inv.compute_handler.get_runtime_key.return_value = 'rk' + meta = inv.select_runtime('M000', 512) + assert meta['lithops_version'] == __version__ + inv.compute_handler.get_runtime_key.assert_called_once() + inv.compute_handler.deploy_runtime.assert_not_called() + # memory override is passed to get_runtime_key + assert inv.compute_handler.get_runtime_key.call_args[0][1] == 512 + + def test_non_serverless_ignores_memory_override(self): + inv = _bare_invoker(mode=LOCALHOST) + inv.internal_storage.get_runtime_meta.return_value = _matching_runtime_meta() + inv.compute_handler.get_runtime_key.return_value = 'rk' + inv.select_runtime('M000', 512) + assert inv.compute_handler.get_runtime_key.call_args[0][1] == 256 + + def test_deploys_runtime_when_meta_missing(self): + inv = _bare_invoker() + inv.internal_storage.get_runtime_meta.return_value = None + inv.compute_handler.get_runtime_key.return_value = 'rk' + inv.compute_handler.deploy_runtime.return_value = _matching_runtime_meta() + meta = inv.select_runtime('M000', None) + inv.compute_handler.deploy_runtime.assert_called_once() + inv.internal_storage.put_runtime_meta.assert_called_once() + assert meta['runtime_timeout'] == 300 + + +class TestPayloadAndFutures: + + def test_create_payload_copies_job_fields(self): + inv = _bare_invoker() + job = _job(chunksize=4, worker_processes=3, extra_env={'K': '1'}) + payload = inv._create_payload(job) + assert payload['func_name'] == 'fn' + assert payload['total_calls'] == 2 + assert payload['call_ids'] is None + assert payload['lithops_version'] == __version__ + assert payload['chunksize'] == 4 + assert payload['worker_processes'] == 3 + assert payload['extra_env'] == {'K': '1'} + assert payload['max_workers'] == 8 + assert payload['data_key'] == 'dk' + + def test_build_futures_marks_invoked_and_copies_metadata(self): + inv = _bare_invoker() + job = _job(total_calls=3, metadata={'func_name': 'fn', 'host_submit_tstamp': 1}) + futures = inv._build_futures(job) + assert len(futures) == 3 + assert job.futures is futures + assert all(isinstance(f, ResponseFuture) for f in futures) + assert all(f.invoked for f in futures) + assert futures[0].call_id == '00000' + assert futures[2].call_id == '00002' + assert futures[0].stats['func_name'] == 'fn' + + def test_run_job_sends_metrics_and_invokes(self, tmp_path, monkeypatch): + monkeypatch.setattr('lithops.invokers.LOGS_DIR', str(tmp_path)) + inv = _bare_invoker() + inv._invoke_job = MagicMock() + job = _job() + futures = inv._run_job(job) + inv._invoke_job.assert_called_once_with(job) + inv.prometheus.send_metric.assert_called() + assert len(futures) == 2 + + def test_run_job_include_function_extends_runtime(self, tmp_path, monkeypatch): + monkeypatch.setattr('lithops.invokers.LOGS_DIR', str(tmp_path)) + inv = _bare_invoker(include_function=True, runtime_name='base:tag') + inv._invoke_job = MagicMock() + with patch('lithops.invokers.extend_runtime') as ext: + job = _job(runtime_name='base:tag') + inv._run_job(job) + ext.assert_called_once() + assert job.runtime_name == 'base:tag' + + def test_run_job_stops_invoker_on_failure(self, tmp_path, monkeypatch): + monkeypatch.setattr('lithops.invokers.LOGS_DIR', str(tmp_path)) + inv = _bare_invoker() + inv._invoke_job = MagicMock(side_effect=RuntimeError('nope')) + inv.stop = MagicMock() + with pytest.raises(RuntimeError, match='nope'): + inv._run_job(_job()) + inv.stop.assert_called_once() + + +class TestBatchInvoker: + + def test_invoke_job_sets_call_ids(self): + inv = BatchInvoker.__new__(BatchInvoker) + for key, value in _bare_invoker().__dict__.items(): + setattr(inv, key, value) + inv.compute_handler.invoke.return_value = 'act-1' + job = _job(total_calls=2) + inv._invoke_job(job) + payload = inv.compute_handler.invoke.call_args[0][0] + assert payload['call_ids'] == ['00000', '00001'] + + def test_run_job_starts_monitor(self, tmp_path, monkeypatch): + monkeypatch.setattr('lithops.invokers.LOGS_DIR', str(tmp_path)) + inv = BatchInvoker.__new__(BatchInvoker) + for key, value in _bare_invoker().__dict__.items(): + setattr(inv, key, value) + inv._invoke_job = MagicMock() + job = _job() + futures = inv.run_job(job) + inv.job_monitor.start.assert_called_once_with(futures) + + +class TestFaaSInvokerHelpers: + + def _faas(self, **attrs): + inv = FaaSInvoker.__new__(FaaSInvoker) + for key, value in _bare_invoker(**attrs).__dict__.items(): + setattr(inv, key, value) + inv.pending_calls_q = queue.Queue() + inv.job_monitor = MagicMock() + inv.job_monitor.token_bucket_q = queue.Queue() + inv.running_workers = 0 + inv.should_run = False + inv.remote_invoker = False + inv.sync = False + inv.invokers = [] + inv.executor = MagicMock() + return inv + + def test_drain_token_bucket_noop_when_no_workers_or_empty(self): + inv = self._faas() + inv.running_workers = 0 + inv._drain_token_bucket() + inv.running_workers = 3 + inv._drain_token_bucket() + assert inv.running_workers == 3 + + def test_drain_token_bucket_consumes_until_zero(self): + inv = self._faas() + inv.running_workers = 2 + inv.job_monitor.token_bucket_q.put('#') + inv.job_monitor.token_bucket_q.put('#') + inv.job_monitor.token_bucket_q.put('#') + inv._drain_token_bucket() + assert inv.running_workers == 0 + assert inv.job_monitor.token_bucket_q.qsize() == 1 + + def test_queue_call_ranges_chunks_ids(self): + inv = self._faas() + job = _job(chunksize=2) + inv._queue_call_ranges(job, range(5)) + ranges = [] + while not inv.pending_calls_q.empty(): + queued_job, ids = inv.pending_calls_q.get() + assert queued_job is job + ranges.append(list(ids)) + assert ranges == [[0, 1], [2, 3], [4]] + + def test_invoke_job_remote_success(self): + inv = self._faas() + inv.compute_handler.invoke.return_value = 'act-r' + inv._invoke_job_remote(_job()) + + def test_invoke_job_remote_failure(self): + inv = self._faas() + inv.compute_handler.invoke.return_value = None + with pytest.raises(Exception, match='Unable to spawn remote invoker'): + inv._invoke_job_remote(_job()) + + def test_invoke_task_requeues_when_activation_missing(self, monkeypatch): + inv = self._faas() + monkeypatch.setattr('lithops.invokers.time.sleep', lambda *_: None) + inv.compute_handler.invoke.return_value = None + job = _job() + inv._invoke_task(job, [0, 1]) + queued_job, ids = inv.pending_calls_q.get_nowait() + assert queued_job is job + assert list(ids) == [0, 1] + assert inv.job_monitor.token_bucket_q.get_nowait() == '#' + + def test_invoke_job_queues_all_when_at_max_workers(self): + inv = self._faas() + inv.should_run = True + inv.running_workers = 8 + inv.max_workers = 8 + job = _job(total_calls=3, chunksize=1) + inv._invoke_job(job) + assert inv.pending_calls_q.qsize() == 3 + + def test_run_job_starts_monitor_with_tokens(self, tmp_path, monkeypatch): + monkeypatch.setattr('lithops.invokers.LOGS_DIR', str(tmp_path)) + inv = self._faas() + inv._invoke_job = MagicMock() + job = _job() + futures = inv.run_job(job) + inv.job_monitor.start.assert_called_once_with( + fs=futures, job_id='M000', chunksize=1, generate_tokens=True + ) + + def test_invoke_job_uses_free_workers_and_queues_remainder(self): + inv = self._faas() + inv.max_workers = 2 + inv.executor = MagicMock() + inv._start_async_invokers = MagicMock() + job = _job(total_calls=5, chunksize=2) + inv._invoke_job(job) + inv._start_async_invokers.assert_called_once() + assert inv.executor.submit.call_count == 2 + assert inv.running_workers == 2 + assert inv.pending_calls_q.qsize() == 1 + _, ids = inv.pending_calls_q.get() + assert list(ids) == [4] + + def test_stop_drains_pending_queue_and_signals_invokers(self): + inv = self._faas() + inv.should_run = True + threads = [MagicMock(), MagicMock()] + inv.invokers = list(threads) + inv.pending_calls_q.put((_job(), [0])) + inv.pending_calls_q.put((_job(), [1])) + inv.stop() + assert inv.should_run is False + assert inv.invokers == [] + for thread in threads: + thread.join.assert_not_called() + tokens = [] + while not inv.job_monitor.token_bucket_q.empty(): + tokens.append(inv.job_monitor.token_bucket_q.get_nowait()) + assert tokens == ['$', '$'] + pending = [] + while not inv.pending_calls_q.empty(): + pending.append(inv.pending_calls_q.get_nowait()) + assert pending == [(None, None), (None, None)] + + def test_start_async_invokers_starts_daemon_threads(self): + inv = self._faas() + inv.invoke_pool_threads = 8 + inv.should_run = True + with patch('lithops.invokers.threading.Thread') as Thread: + inv._start_async_invokers() + assert Thread.call_count == FaaSInvoker.ASYNC_INVOKERS + thread = Thread.return_value + assert thread.daemon is True + assert thread.start.call_count == FaaSInvoker.ASYNC_INVOKERS + assert inv.job_monitor.token_bucket_q.qsize() == FaaSInvoker.ASYNC_INVOKERS + assert len(inv.invokers) == FaaSInvoker.ASYNC_INVOKERS + + def test_invoke_task_uses_data_byte_strs_when_no_data_key(self): + inv = self._faas() + inv.compute_handler.invoke.return_value = 'act-1' + job = _job(data_key=None, data_byte_strs=[b'a', b'b'], total_calls=2) + inv._invoke_task(job, [0, 1]) + payload = inv.compute_handler.invoke.call_args[0][0] + assert 'data_byte_ranges' not in payload + assert payload['data_byte_strs'] == [b'a', b'b'] + assert payload['call_ids'] == ['00000', '00001'] + + def test_invoke_task_slices_data_byte_ranges_for_chunk(self): + inv = self._faas() + inv.compute_handler.invoke.return_value = 'act-1' + job = _job( + total_calls=4, + data_byte_ranges=[(0, 1), (2, 3), (4, 5), (6, 7)], + ) + inv._invoke_task(job, [1, 2]) + payload = inv.compute_handler.invoke.call_args[0][0] + assert payload['call_ids'] == ['00001', '00002'] + assert payload['data_byte_ranges'] == [(2, 3), (4, 5)] + assert payload['chunksize'] == 1 + + def test_invoke_job_remote_flag_skips_local_scheduling(self): + inv = self._faas() + inv.remote_invoker = True + inv.compute_handler.invoke.return_value = 'act-r' + inv._start_async_invokers = MagicMock() + inv._invoke_job(_job()) + inv._start_async_invokers.assert_not_called() + inv.compute_handler.pre_invoke.assert_called_once() + payload = inv.compute_handler.invoke.call_args[0][0] + assert payload['remote_invoker'] is True + assert payload['job']['job_id'] == 'M000' + assert inv.pending_calls_q.empty() + + def test_invoke_job_exact_fit_does_not_queue(self): + inv = self._faas() + inv.max_workers = 3 + inv.executor = MagicMock() + inv._start_async_invokers = MagicMock() + inv._invoke_job(_job(total_calls=3, chunksize=1)) + assert inv.executor.submit.call_count == 3 + assert inv.pending_calls_q.empty() + assert inv.running_workers == 3 + + def test_second_job_at_capacity_queues_all_without_restart(self): + inv = self._faas() + inv.should_run = True + inv.running_workers = 8 + inv.max_workers = 8 + inv._start_async_invokers = MagicMock() + inv._invoke_job(_job(total_calls=3, chunksize=1)) + inv._start_async_invokers.assert_not_called() + assert inv.pending_calls_q.qsize() == 3 + + def test_second_job_drains_leftover_tokens_before_direct_invoke(self): + inv = self._faas() + inv.should_run = True + inv.running_workers = 3 + inv.max_workers = 8 + inv.executor = MagicMock() + inv._start_async_invokers = MagicMock() + inv.job_monitor.token_bucket_q.put('#') + inv.job_monitor.token_bucket_q.put('#') + inv._invoke_job(_job(total_calls=2, chunksize=1)) + inv._start_async_invokers.assert_not_called() + assert inv.job_monitor.token_bucket_q.empty() + assert inv.executor.submit.call_count == 2 + assert inv.running_workers == 3 + assert inv.pending_calls_q.empty() + + def test_stop_waits_for_invoker_threads_when_asked(self): + inv = self._faas() + inv.should_run = True + threads = [MagicMock(), MagicMock()] + inv.invokers = list(threads) + inv.stop(wait=True) + for thread in threads: + thread.join.assert_called_once_with(timeout=inv.STOP_TIMEOUT) + + def test_stop_with_wait_blocks_until_the_invocations_are_done(self): + # This is what replaced the blind sleep(5) in the remote invoker: the + # async invoker threads drain the invocations already in flight only + # after they leave their loop, and nothing else joins them + inv = self._faas() + inv.should_run = True + started = threading.Event() + finished = [] + + def in_flight(): + started.set() + time.sleep(0.3) + finished.append('done') + + thread = threading.Thread(target=in_flight) + inv.invokers = [thread] + thread.start() + assert started.wait(timeout=5) + + inv.stop(wait=True) + assert finished == ['done'], 'stop returned before the call finished' + assert not thread.is_alive() + + def test_stop_without_wait_returns_while_calls_are_in_flight(self): + inv = self._faas() + inv.should_run = True + started = threading.Event() + release = threading.Event() + + def in_flight(): + started.set() + release.wait(timeout=5) + + thread = threading.Thread(target=in_flight, daemon=True) + inv.invokers = [thread] + thread.start() + assert started.wait(timeout=5) + try: + inv.stop() + assert thread.is_alive(), 'stop should not have waited' + finally: + release.set() + thread.join(timeout=5) + + def test_stop_is_noop_when_no_async_invokers(self): + inv = self._faas() + inv.should_run = True + inv.pending_calls_q.put((_job(), [0])) + inv.stop() + assert inv.should_run is True + assert inv.pending_calls_q.qsize() == 1 + + def test_invoke_job_calls_pre_invoke(self): + inv = self._faas() + inv.should_run = True + inv.running_workers = 8 + inv.max_workers = 8 + job = _job(total_calls=1, chunksize=1) + inv._invoke_job(job) + inv.compute_handler.pre_invoke.assert_called_once_with(job) + + +class TestFaaSInvokerInit: + + def _handler(self): + handler = MagicMock() + handler.get_runtime_info.return_value = { + 'runtime_name': 'python:3', + 'runtime_memory': 256, + 'runtime_timeout': 300, + 'max_workers': 8, + } + return handler + + def _config(self, **backend): + ibm_cf = { + 'invoke_pool_threads': 4, + 'remote_invoker': False, + } + ibm_cf.update(backend) + return { + 'lithops': { + 'mode': SERVERLESS, + 'backend': 'ibm_cf', + 'storage': 'localhost', + 'telemetry': False, + }, + 'ibm_cf': ibm_cf, + 'localhost': {'storage_bucket': 'test-bucket'}, + } + + def test_init_reads_pool_threads_and_remote_invoker(self): + inv = FaaSInvoker( + self._config(remote_invoker=True, invoke_pool_threads=6), + 'sess-0', + MagicMock(), + self._handler(), + MagicMock(), + ) + try: + assert inv.remote_invoker is True + assert inv.sync is False + assert inv.max_workers == 8 + assert inv.invoke_pool_threads == 6 + assert inv.should_run is False + assert inv.pending_calls_q.empty() + finally: + inv.executor.shutdown(wait=False) + + def test_init_disables_remote_invoker_inside_worker(self, monkeypatch): + monkeypatch.setenv('LITHOPS_WORKER', '1') + inv = FaaSInvoker( + self._config(remote_invoker=True), + 'sess-0', + MagicMock(), + self._handler(), + MagicMock(), + ) + try: + assert inv.remote_invoker is False + assert inv.sync is True + finally: + inv.executor.shutdown(wait=False) + + +def _wait_until(predicate, timeout=5): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return + time.sleep(0.01) + raise AssertionError(f'timed out after {timeout}s') + + +class TestFaaSTokenBucketLoop: + """Closed-loop token bucket with real invoker threads and a mocked backend.""" + + def _live_faas(self, *, max_workers=1, async_invokers=1, pool_workers=16): + inv = FaaSInvoker.__new__(FaaSInvoker) + for key, value in _bare_invoker().__dict__.items(): + setattr(inv, key, value) + inv.pending_calls_q = queue.Queue() + inv.job_monitor = MagicMock() + inv.job_monitor.token_bucket_q = queue.Queue() + inv.running_workers = 0 + inv.should_run = False + inv.remote_invoker = False + inv.sync = False + inv.invokers = [] + inv.executor = ThreadPoolExecutor(max_workers=pool_workers) + inv.invoke_pool_threads = max(8, pool_workers) + inv.ASYNC_INVOKERS = async_invokers + inv.max_workers = max_workers + return inv + + def test_token_unblocks_queued_invocation(self): + invoked = [] + lock = threading.Lock() + + def invoke(payload): + with lock: + invoked.append(list(payload['call_ids'])) + return 'act-1' + + inv = self._live_faas() + inv.compute_handler.invoke.side_effect = invoke + inv.should_run = True + job = _job(total_calls=2, chunksize=1) + try: + inv._start_async_invokers() + inv._queue_call_ranges(job, range(2)) + _wait_until(lambda: len(invoked) == 1) + with lock: + assert invoked == [['00000']] + assert inv.pending_calls_q.qsize() == 1 + + inv.job_monitor.token_bucket_q.put('#') + _wait_until(lambda: len(invoked) == 2) + with lock: + assert invoked == [['00000'], ['00001']] + finally: + inv.stop() + inv.executor.shutdown(wait=True) + + def test_overflow_beyond_max_workers_waits_for_tokens(self): + """10 workers, 20 calls: first wave runs now, the rest wait for tokens.""" + invoked = [] + lock = threading.Lock() + + def invoke(payload): + with lock: + invoked.append(list(payload['call_ids'])) + return 'act-1' + + max_workers = 10 + total_calls = 20 + async_invokers = FaaSInvoker.ASYNC_INVOKERS + first_wave = max_workers + async_invokers + leftover = total_calls - first_wave + + inv = self._live_faas( + max_workers=max_workers, async_invokers=async_invokers + ) + inv.compute_handler.invoke.side_effect = invoke + job = _job( + total_calls=total_calls, + chunksize=1, + data_byte_ranges=[(i, i) for i in range(total_calls)], + ) + try: + inv._invoke_job(job) + _wait_until(lambda: len(invoked) == first_wave) + assert inv.pending_calls_q.qsize() == leftover + with lock: + ids = [call_id for chunk in invoked for call_id in chunk] + assert len(set(ids)) == first_wave + + for _ in range(leftover): + inv.job_monitor.token_bucket_q.put('#') + _wait_until(lambda: len(invoked) == total_calls) + with lock: + ids = [call_id for chunk in invoked for call_id in chunk] + assert sorted(ids) == [f'{i:05d}' for i in range(total_calls)] + finally: + inv.stop() + inv.executor.shutdown(wait=True) + + def test_completions_refill_bucket_and_drain_overflow(self): + """Finished workers put tokens so the queued extra 10 calls all run.""" + invoked = [] + lock = threading.Lock() + inv = self._live_faas( + max_workers=10, async_invokers=FaaSInvoker.ASYNC_INVOKERS + ) + + def invoke(payload): + with lock: + invoked.append(list(payload['call_ids'])) + inv.job_monitor.token_bucket_q.put('#') + return 'act-1' + + inv.compute_handler.invoke.side_effect = invoke + job = _job( + total_calls=20, + chunksize=1, + data_byte_ranges=[(i, i) for i in range(20)], + ) + try: + inv._invoke_job(job) + _wait_until(lambda: len(invoked) == 20) + with lock: + ids = [call_id for chunk in invoked for call_id in chunk] + assert sorted(ids) == [f'{i:05d}' for i in range(20)] + finally: + inv.stop() + inv.executor.shutdown(wait=True) + + def test_overflow_with_chunksize_releases_one_token_per_worker(self): + invoked = [] + lock = threading.Lock() + + def invoke(payload): + with lock: + invoked.append(list(payload['call_ids'])) + return 'act-1' + + inv = self._live_faas(max_workers=2, async_invokers=1) + inv.compute_handler.invoke.side_effect = invoke + job = _job( + total_calls=8, + chunksize=2, + data_byte_ranges=[(i, i) for i in range(8)], + ) + try: + inv._invoke_job(job) + # 2 direct workers + 1 primed async worker = 3 invokes, 6 calls + _wait_until(lambda: len(invoked) == 3) + assert inv.pending_calls_q.qsize() == 1 + inv.job_monitor.token_bucket_q.put('#') + _wait_until(lambda: len(invoked) == 4) + with lock: + ids = [call_id for chunk in invoked for call_id in chunk] + assert sorted(ids) == [f'{i:05d}' for i in range(8)] + finally: + inv.stop() + inv.executor.shutdown(wait=True) + + def test_sync_waits_until_direct_invokes_finish(self): + started = threading.Event() + release = threading.Event() + + def invoke(payload): + started.set() + assert release.wait(timeout=2) + return 'act-1' + + inv = self._live_faas(max_workers=1, async_invokers=1) + inv.sync = True + inv.compute_handler.invoke.side_effect = invoke + job = _job(total_calls=1, chunksize=1) + try: + worker = threading.Thread(target=inv._invoke_job, args=(job,)) + worker.start() + assert started.wait(timeout=2) + assert worker.is_alive() + release.set() + worker.join(timeout=2) + assert not worker.is_alive() + assert inv.compute_handler.invoke.call_count == 1 + finally: + release.set() + inv.stop() + inv.executor.shutdown(wait=True) + + def test_second_job_queues_until_tokens_after_workers_are_full(self): + invoked = [] + lock = threading.Lock() + + def invoke(payload): + with lock: + invoked.append(list(payload['call_ids'])) + return 'act-1' + + inv = self._live_faas(max_workers=2, async_invokers=1) + inv.compute_handler.invoke.side_effect = invoke + job1 = _job( + total_calls=2, + chunksize=1, + data_byte_ranges=[(0, 0), (1, 1)], + ) + job2 = _job( + job_id='M001', + total_calls=2, + chunksize=1, + data_byte_ranges=[(0, 0), (1, 1)], + ) + try: + inv._invoke_job(job1) + _wait_until(lambda: len(invoked) == 2) + inv._invoke_job(job2) + assert inv.pending_calls_q.qsize() == 2 + with lock: + assert len(invoked) == 2 + inv.job_monitor.token_bucket_q.put('#') + inv.job_monitor.token_bucket_q.put('#') + _wait_until(lambda: len(invoked) == 4) + finally: + inv.stop() + inv.executor.shutdown(wait=True) + + def test_run_job_invokes_and_starts_monitor_with_tokens(self, tmp_path, monkeypatch): + monkeypatch.setattr('lithops.invokers.LOGS_DIR', str(tmp_path)) + invoked = [] + lock = threading.Lock() + + def invoke(payload): + with lock: + invoked.append(list(payload['call_ids'])) + return 'act-1' + + inv = self._live_faas(max_workers=2, async_invokers=1) + inv.compute_handler.invoke.side_effect = invoke + job = _job( + total_calls=2, + chunksize=1, + data_byte_ranges=[(0, 0), (1, 1)], + ) + try: + futures = inv.run_job(job) + _wait_until(lambda: len(invoked) == 2) + assert len(futures) == 2 + assert all(f.invoked for f in futures) + inv.job_monitor.start.assert_called_once_with( + fs=futures, + job_id='M000', + chunksize=1, + generate_tokens=True, + ) + finally: + inv.stop() + inv.executor.shutdown(wait=True) + + def test_failed_invoke_returns_token_and_retries(self, monkeypatch): + monkeypatch.setattr('lithops.invokers.time.sleep', lambda *_: None) + invoked = [] + lock = threading.Lock() + attempts = {'n': 0} + + def invoke(payload): + with lock: + invoked.append(list(payload['call_ids'])) + attempts['n'] += 1 + if attempts['n'] == 1: + return None + return 'act-1' + + inv = self._live_faas() + inv.max_workers = 0 + inv.compute_handler.invoke.side_effect = invoke + job = _job(total_calls=1, chunksize=1) + try: + inv._invoke_job(job) + _wait_until(lambda: len(invoked) >= 2) + with lock: + assert invoked[0] == ['00000'] + assert invoked[1] == ['00000'] + finally: + inv.stop() + inv.executor.shutdown(wait=True) + + +class TestExtendRuntime: + + def test_skips_build_when_meta_already_exists(self): + job = SimpleNamespace( + runtime_name='img:tag', + ext_runtime_uuid='abc123', + runtime_memory=256, + runtime_timeout=60, + ) + compute = MagicMock() + compute.get_runtime_key.return_value = 'rk' + internal = MagicMock() + internal.get_runtime_meta.return_value = _matching_runtime_meta() + extend_runtime(job, compute, internal) + assert job.runtime_name == 'img:abc123' + compute.build_runtime.assert_not_called() + compute.deploy_runtime.assert_not_called() + + def test_builds_and_deploys_when_meta_missing(self, tmp_path, monkeypatch): + local = tmp_path / 'ext' + local.mkdir() + job = SimpleNamespace( + runtime_name='img:tag', + ext_runtime_uuid='abc123', + runtime_memory=256, + runtime_timeout=60, + local_tmp_dir=str(local), + ) + compute = MagicMock() + compute.get_runtime_key.return_value = 'rk' + compute.deploy_runtime.return_value = _matching_runtime_meta() + internal = MagicMock() + internal.get_runtime_meta.return_value = None + monkeypatch.chdir(tmp_path) + extend_runtime(job, compute, internal) + compute.build_runtime.assert_called_once() + compute.deploy_runtime.assert_called_once() + internal.put_runtime_meta.assert_called_once() + assert not local.exists() + assert job.runtime_name == 'img:abc123' + + def test_restores_cwd_if_build_runtime_raises(self, tmp_path, monkeypatch): + local = tmp_path / 'ext' + local.mkdir() + job = SimpleNamespace( + runtime_name='img:tag', + ext_runtime_uuid='abc123', + runtime_memory=256, + runtime_timeout=60, + local_tmp_dir=str(local), + ) + compute = MagicMock() + compute.get_runtime_key.return_value = 'rk' + compute.build_runtime.side_effect = RuntimeError('build failed') + internal = MagicMock() + internal.get_runtime_meta.return_value = None + monkeypatch.chdir(tmp_path) + cwd = os.getcwd() + with pytest.raises(RuntimeError, match='build failed'): + extend_runtime(job, compute, internal) + assert os.getcwd() == cwd diff --git a/lithops/tests/test_job.py b/lithops/tests/test_job.py new file mode 100644 index 000000000..44b155352 --- /dev/null +++ b/lithops/tests/test_job.py @@ -0,0 +1,1178 @@ +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import errno +import logging +import os +import pickle +from functools import partial +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from lithops.constants import LOCALHOST, MAX_AGG_DATA_SIZE, SERVERLESS, STANDALONE +from lithops.job import create_map_job, create_reduce_job +from lithops.job.job import ( + FUNCTION_CACHE, + MAX_DATA_IN_PAYLOAD, + _FUNC_SERIALIZE_CACHE, + _store_func_and_modules, + invalidate_function_cache, +) +from lithops.job.partitioner import CHUNK_THRESHOLD, create_partitions +from lithops.job.serialize import SerializeIndependent, create_module_data +from lithops.storage.utils import ( + CloudObject, + CloudObjectLocal, + CloudObjectUrl, + create_func_key, + func_key_suffix, +) +from lithops.utils import BackendType, bytes_to_b64str + + +def _echo(x): + return x + + +def _obj_fn(obj): + return obj + + +def _reduce_fn(results): + return results + + +class _Adder: + def __call__(self, x): + return x + 1 + + +class _NoCall: + pass + + +class _CapturingSerializer: + last = None + + def __init__(self, preinstalls): + self.preinstalls = preinstalls + + def __call__(self, objs, inc_modules, exc_modules): + type(self).last = (inc_modules, exc_modules, objs) + return ([b'func'] + [b'data'] * (len(objs) - 1), set()) + + +def _job_config( + mode=LOCALHOST, + backend=None, + backend_type=BackendType.FAAS.value, + **lithops +): + backend = backend or mode + return { + 'lithops': { + 'mode': mode, + 'backend': backend, + 'chunksize': 1, + 'execution_timeout': 100, + 'backend_type': backend_type, + **lithops, + }, + backend: { + 'worker_processes': 2, + 'runtime_memory': 256, + }, + STANDALONE: { + 'worker_processes': 2, + 'hard_dismantle_timeout': 50, + }, + } + + +def _runtime_meta(**extra): + meta = {'preinstalls': [['os', True]], 'runtime_timeout': 30} + meta.update(extra) + return meta + + +def _storage(): + internal = MagicMock() + internal.backend = 'localhost' + internal.storage = MagicMock() + return internal + + +@pytest.fixture +def fresh_function_cache(): + saved = set(FUNCTION_CACHE) + saved_serialize = {func: dict(entries) for func, entries in _FUNC_SERIALIZE_CACHE.items()} + FUNCTION_CACHE.clear() + _FUNC_SERIALIZE_CACHE.clear() + yield FUNCTION_CACHE + FUNCTION_CACHE.clear() + FUNCTION_CACHE.update(saved) + _FUNC_SERIALIZE_CACHE.clear() + _FUNC_SERIALIZE_CACHE.update(saved_serialize) + + +@pytest.fixture +def capturing_serializer(monkeypatch): + _CapturingSerializer.last = None + monkeypatch.setattr( + 'lithops.job.job.SerializeIndependent', _CapturingSerializer + ) + return _CapturingSerializer + + +def _make_map_job( + *, + func=_echo, + iterdata=[1, 2], + extra_env=None, + include_modules=[], + exclude_modules=None, + execution_timeout=None, + runtime_memory=None, + chunksize=None, + extra_args=None, + config=None, + internal_storage=None, + executor_id='exec', + job_id='j0', + runtime_meta=None, + **map_kwargs +): + return create_map_job( + config=config or _job_config(), + internal_storage=internal_storage or _storage(), + executor_id=executor_id, + job_id=job_id, + map_function=func, + iterdata=iterdata, + runtime_meta=runtime_meta or _runtime_meta(), + runtime_memory=runtime_memory, + extra_env=extra_env, + include_modules=include_modules, + exclude_modules=exclude_modules, + execution_timeout=execution_timeout, + chunksize=chunksize, + extra_args=extra_args, + **map_kwargs + ) + + +@pytest.mark.usefixtures('capturing_serializer', 'fresh_function_cache') +class TestCreateMapJob: + + def test_basic_job_fields(self): + storage = _storage() + job = _make_map_job(internal_storage=storage) + assert job.executor_id == 'exec' + assert job.job_id == 'j0' + assert job.job_key == 'exec-j0' + assert job.function_name == '_echo' + assert job.total_calls == 2 + assert job.chunksize == 1 + assert job.worker_processes == 2 + assert job.runtime_memory is None + assert job.runtime_timeout is None + assert job.data_key is None + assert job.data_byte_ranges is None + assert job.data_byte_strs == [b'data', b'data'] + storage.put_func.assert_called_once() + storage.put_data.assert_not_called() + assert job.func_key in FUNCTION_CACHE + + def test_callable_class_function_name(self): + job = _make_map_job(func=_Adder()) + assert job.function_name == '_Adder' + + def test_bound_method_function_name(self): + job = _make_map_job(func=_Adder().__call__) + assert job.function_name == '__call__' + + def test_extra_env_bools_converted_on_copy(self): + extra = {'FLAG': True, 'N': 1} + job = _make_map_job(extra_env=extra) + assert extra['FLAG'] is True + assert job.extra_env['FLAG'] == 'True' + assert job.extra_env['N'] == 1 + + def test_none_extra_env_becomes_empty(self): + job = _make_map_job() + assert job.extra_env == {} + + def test_zero_timeout_and_chunksize_are_kept(self): + job = _make_map_job(execution_timeout=0, chunksize=0) + assert job.execution_timeout == 0 + assert job.chunksize == 0 + + def test_chunksize_overrides_config_default(self): + job = _make_map_job(chunksize=7) + assert job.chunksize == 7 + + def test_serverless_clamps_timeout(self): + cfg = _job_config(mode=SERVERLESS, backend='ibm_cf') + job = _make_map_job(config=cfg, execution_timeout=100) + assert job.runtime_memory == 256 + assert job.runtime_timeout == 30 + assert job.execution_timeout == 25 + + def test_serverless_keeps_timeout_below_runtime(self): + cfg = _job_config(mode=SERVERLESS, backend='ibm_cf') + job = _make_map_job(config=cfg, execution_timeout=10, runtime_memory=512) + assert job.runtime_memory == 512 + assert job.execution_timeout == 10 + + def test_standalone_clamps_timeout(self): + cfg = _job_config(mode=STANDALONE, backend='aws_ec2') + job = _make_map_job(config=cfg) + assert job.runtime_memory is None + assert job.execution_timeout == 40 + + def test_unknown_mode_is_not_standalone_substring(self): + cfg = _job_config(mode='a', backend='aws_ec2') + job = _make_map_job(config=cfg) + assert job.runtime_memory is None + assert job.runtime_timeout is None + assert job.execution_timeout == 100 + + def test_function_cache_skips_second_upload(self): + storage = _storage() + _make_map_job(internal_storage=storage) + _make_map_job(internal_storage=storage) + storage.put_func.assert_called_once() + + def test_function_cache_reuploads_after_invalidate(self): + storage = _storage() + job = _make_map_job(internal_storage=storage) + storage.put_func.assert_called_once() + invalidate_function_cache(job.executor_id) + _make_map_job(internal_storage=storage) + assert storage.put_func.call_count == 2 + + def test_batch_backend_always_uploads_data(self): + storage = _storage() + cfg = _job_config(backend_type=BackendType.BATCH.value) + job = _make_map_job(config=cfg, internal_storage=storage) + storage.put_data.assert_called_once() + assert job.data_key is not None + assert job.data_byte_ranges is not None + assert not hasattr(job, 'data_byte_strs') + + def test_data_limit_raises(self): + cfg = _job_config(data_limit=1e-9) + with pytest.raises(Exception, match='exceeded maximum size'): + _make_map_job(config=cfg) + + def test_data_limit_zero_skips_check(self): + job = _make_map_job(config=_job_config(data_limit=0)) + assert job.total_calls == 2 + + def test_missing_data_limit_uses_constant(self, monkeypatch): + monkeypatch.setattr('lithops.job.job.MAX_AGG_DATA_SIZE', 1e-12) + cfg = _job_config() + assert 'data_limit' not in cfg['lithops'] + with pytest.raises(Exception, match='exceeded maximum size'): + _make_map_job(config=cfg) + assert MAX_AGG_DATA_SIZE == 4 + + def test_include_modules_none_string(self): + _make_map_job(config=_job_config(include_modules='none')) + inc, exc, _ = _CapturingSerializer.last + assert inc is None + assert exc == set() + + def test_include_modules_invalid_string_raises(self): + cfg = _job_config(include_modules='all') + with pytest.raises(ValueError, match='must be a list'): + _make_map_job(config=cfg) + + def test_include_modules_cfg_none_and_arg_none(self): + _make_map_job( + config=_job_config(include_modules='NONE'), + include_modules=None, + ) + inc, _, _ = _CapturingSerializer.last + assert inc is None + + def test_include_modules_empty_cfg_and_empty_arg_is_empty_set(self): + _make_map_job(include_modules=[]) + inc, _, _ = _CapturingSerializer.last + assert inc == set() + + def test_include_modules_arg_none_overrides_empty_cfg(self): + _make_map_job(include_modules=None) + inc, _, _ = _CapturingSerializer.last + assert inc is None + + def test_include_and_exclude_union(self): + cfg = _job_config(include_modules=['a'], exclude_modules=['x']) + _make_map_job(config=cfg, include_modules=['b'], exclude_modules=['y']) + inc, exc, _ = _CapturingSerializer.last + assert inc == {'a', 'b'} + assert exc == {'x', 'y'} + + def test_object_processing_sets_parts(self): + with patch( + 'lithops.job.job.create_partitions', + return_value=([{'obj': 'x'}], [2, 3]), + ) as cp: + job = _make_map_job(func=_obj_fn, iterdata=['http://example.com/a']) + assert job.parts_per_object == [2, 3] + assert cp.called + assert 'host_job_create_partitions_time' in job.metadata + + def test_empty_parts_per_object_not_attached(self): + with patch( + 'lithops.job.job.create_partitions', + return_value=([{'obj': 'x'}], []), + ): + job = _make_map_job(func=_obj_fn, iterdata=['http://example.com/a']) + assert not hasattr(job, 'parts_per_object') + + def test_empty_object_iterdata_creates_empty_job(self): + job = _make_map_job(func=_obj_fn, iterdata=[]) + assert job.total_calls == 0 + assert not hasattr(job, 'parts_per_object') + + def test_runtime_include_function(self, tmp_path, monkeypatch): + monkeypatch.setattr('lithops.job.job.CUSTOM_RUNTIME_DIR', str(tmp_path)) + cfg = _job_config() + cfg[LOCALHOST]['runtime_include_function'] = True + storage = _storage() + job = _make_map_job(config=cfg, internal_storage=storage) + storage.put_func.assert_not_called() + assert job.func_key == func_key_suffix + assert job.ext_runtime_uuid + assert os.path.isdir(job.local_tmp_dir) + func_path = os.path.join(job.local_tmp_dir, job.func_key) + assert os.path.isfile(func_path) + with open(func_path, 'rb') as f: + assert pickle.load(f) == {'func': b'func'} + + def test_metadata_timings_present(self): + job = _make_map_job() + meta = job.metadata + assert 'host_job_create_tstamp' in meta + assert 'host_job_serialize_time' in meta + assert 'func_data_size_bytes' in meta + assert 'func_module_size_bytes' in meta + assert 'host_func_upload_time' in meta + assert 'host_data_upload_time' in meta + assert 'host_job_created_time' in meta + + +@pytest.mark.usefixtures('capturing_serializer', 'fresh_function_cache') +class TestCreateReduceJob: + + def test_single_iterdata_without_parts(self): + job = create_reduce_job( + config=_job_config(), + internal_storage=_storage(), + executor_id='exec', + reduce_job_id='r0', + reduce_function=_reduce_fn, + map_job=SimpleNamespace(), + map_futures=list(range(6)), + runtime_meta=_runtime_meta(), + runtime_memory=None, + obj_reduce_by_key=False, + extra_env=None, + include_modules=[], + exclude_modules=None, + ) + assert job.job_id == 'r0' + assert job.total_calls == 1 + assert job.extra_env['__LITHOPS_REDUCE_JOB'] == 'True' + _, _, objs = _CapturingSerializer.last + assert objs[1]['results'] == list(range(6)) + + def test_reduce_by_key_slices_futures(self): + job = create_reduce_job( + config=_job_config(), + internal_storage=_storage(), + executor_id='exec', + reduce_job_id='r0', + reduce_function=_reduce_fn, + map_job=SimpleNamespace(parts_per_object=[2, 3, 1]), + map_futures=['a', 'b', 'c', 'd', 'e', 'f'], + runtime_meta=_runtime_meta(), + runtime_memory=None, + obj_reduce_by_key=True, + extra_env={'K': False}, + include_modules=[], + exclude_modules=None, + ) + assert job.total_calls == 3 + assert job.extra_env['K'] == 'False' + assert job.extra_env['__LITHOPS_REDUCE_JOB'] == 'True' + _, _, objs = _CapturingSerializer.last + assert [o['results'] for o in objs[1:]] == [ + ['a', 'b'], ['c', 'd', 'e'], ['f'] + ] + + def test_parts_without_reduce_by_key_stays_one_call(self): + job = create_reduce_job( + config=_job_config(), + internal_storage=_storage(), + executor_id='exec', + reduce_job_id='r0', + reduce_function=_reduce_fn, + map_job=SimpleNamespace(parts_per_object=[2, 2]), + map_futures=list(range(4)), + runtime_meta=_runtime_meta(), + runtime_memory=None, + obj_reduce_by_key=False, + extra_env=None, + include_modules=[], + exclude_modules=None, + ) + assert job.total_calls == 1 + + +class TestLargePayloadAndCache: + + def test_large_payload_uploads_data(self, monkeypatch, fresh_function_cache): + class _Big: + def __init__(self, preinstalls): + pass + + def __call__(self, objs, inc, exc): + payload = b'x' * (MAX_DATA_IN_PAYLOAD + 1) + return ([b'f'] + [payload] * (len(objs) - 1), set()) + + monkeypatch.setattr('lithops.job.job.SerializeIndependent', _Big) + storage = _storage() + job = create_map_job( + config=_job_config(), + internal_storage=storage, + executor_id='exec', + job_id='j0', + map_function=_echo, + iterdata=[1], + runtime_meta=_runtime_meta(), + runtime_memory=None, + extra_env=None, + include_modules=[], + exclude_modules=None, + execution_timeout=None, + ) + storage.put_data.assert_called_once() + assert job.data_key is not None + + +class TestFuncSerializeCache: + + def test_second_map_skips_cloudpickle_of_same_function( + self, monkeypatch, fresh_function_cache + ): + dumped = [] + import cloudpickle + real_dumps = cloudpickle.dumps + + def tracking_dumps(obj): + dumped.append(obj) + return real_dumps(obj) + + monkeypatch.setattr( + 'lithops.job.serialize.cloudpickle.dumps', tracking_dumps + ) + storage = _storage() + _make_map_job( + internal_storage=storage, func=_echo, iterdata=[1, 2] + ) + assert dumped.count(_echo) == 1 + data_dumps = [obj for obj in dumped if obj is not _echo] + assert len(data_dumps) == 2 + + _make_map_job( + internal_storage=storage, func=_echo, iterdata=[3, 4] + ) + assert dumped.count(_echo) == 1 + data_dumps = [obj for obj in dumped if obj is not _echo] + assert len(data_dumps) == 4 + + +class TestInvalidateFunctionCache: + + def test_drops_only_matching_executor(self, fresh_function_cache): + keep = create_func_key('other-1', 'aaa') + drop = create_func_key('exec-0', 'bbb') + similar = create_func_key('exec-0-extra', 'ccc') + FUNCTION_CACHE.update({keep, drop, similar}) + invalidate_function_cache('exec-0') + assert keep in FUNCTION_CACHE + assert drop not in FUNCTION_CACHE + assert similar in FUNCTION_CACHE + + +class TestStoreFuncAndModules: + + def test_writes_func_pickle(self, tmp_path): + _store_func_and_modules(str(tmp_path), 'func.pickle', b'abc', None) + with open(tmp_path / 'func.pickle', 'rb') as f: + assert pickle.load(f) == {'func': b'abc'} + + def test_writes_modules_and_strips_leading_slash(self, tmp_path): + payload = bytes_to_b64str(b'hello') + _store_func_and_modules( + str(tmp_path), + 'func.pickle', + b'f', + {'/pkg/mod.py': payload, '/pkg/other.py': payload}, + ) + assert (tmp_path / 'modules' / 'pkg' / 'mod.py').read_bytes() == b'hello' + assert (tmp_path / 'modules' / 'pkg' / 'other.py').read_bytes() == b'hello' + + def test_writes_nested_posix_module_keys_to_native_paths(self, tmp_path): + payload = bytes_to_b64str(b'nested') + _store_func_and_modules( + str(tmp_path), + 'func.pickle', + b'f', + {'pkg/sub/mod.py': payload}, + ) + written = tmp_path / 'modules' / 'pkg' / 'sub' / 'mod.py' + assert written.read_bytes() == b'nested' + + def test_makedirs_errno_17_is_ignored(self, tmp_path): + payload = bytes_to_b64str(b'x') + _store_func_and_modules( + str(tmp_path), + 'func.pickle', + b'f', + {'pkg/a.py': payload, 'pkg/b.py': payload}, + ) + assert (tmp_path / 'modules' / 'pkg' / 'a.py').read_bytes() == b'x' + assert (tmp_path / 'modules' / 'pkg' / 'b.py').read_bytes() == b'x' + + def test_makedirs_uses_exist_ok(self, tmp_path, monkeypatch): + real = os.makedirs + + def _makedirs(path, exist_ok=False): + if not exist_ok: + raise OSError(errno.EACCES, 'denied') + return real(path, exist_ok=exist_ok) + + monkeypatch.setattr(os, 'makedirs', _makedirs) + _store_func_and_modules( + str(tmp_path), + 'func.pickle', + b'f', + {'a.py': bytes_to_b64str(b'x')}, + ) + assert (tmp_path / 'modules' / 'a.py').read_bytes() == b'x' + + +class TestSerializeIndependent: + + def test_init_does_not_mutate_preinstalls(self): + pre = [['os', True]] + ser = SerializeIndependent(pre) + assert pre == [['os', True]] + assert ser.preinstalled_modules == [['os', True], ['lithops', True]] + SerializeIndependent(pre) + assert pre == [['os', True]] + + def test_module_paths_does_not_dump(self, monkeypatch): + dumped = [] + import cloudpickle + real_dumps = cloudpickle.dumps + + def tracking_dumps(obj): + dumped.append(obj) + return real_dumps(obj) + + monkeypatch.setattr( + 'lithops.job.serialize.cloudpickle.dumps', tracking_dumps + ) + ser = SerializeIndependent([['os', True]]) + paths = ser.module_paths([_echo], None, set()) + assert paths == set() + assert dumped == [] + strs = ser.dumps([_echo, 1]) + assert dumped == [_echo, 1] + assert pickle.loads(strs[1]) == 1 + + def test_include_modules_none_skips_manager(self): + ser = SerializeIndependent([['os', True]]) + strs, paths = ser([_echo, 1], None, []) + assert len(strs) == 2 + assert paths == set() + assert pickle.loads(strs[1]) == 1 + + def test_explicit_include_skips_missing_file(self): + ser = SerializeIndependent([['os', True]]) + _, paths = ser([_echo], ['no-such-module.py'], []) + assert paths == set() + + def test_explicit_include_existing_py_file(self, tmp_path): + mod = tmp_path / 'custom.py' + mod.write_text('x = 1\n') + ser = SerializeIndependent([['os', True]]) + _, paths = ser([_echo], [str(mod)], []) + assert os.path.abspath(str(mod)) in paths + + def test_explicit_include_skips_preinstalled(self): + ser = SerializeIndependent([['json', True]]) + _, paths = ser([_echo], ['json.decoder'], []) + assert paths == set() + + def test_explicit_include_importable_module(self): + ser = SerializeIndependent([['os', True]]) + _, paths = ser([_echo], ['json'], []) + assert paths + assert any('json' in p for p in paths) + + def test_explicit_include_missing_module(self): + ser = SerializeIndependent([['os', True]]) + _, paths = ser([_echo], ['definitely_not_a_module_xyz'], []) + assert paths == set() + + def test_class_without_call_raises(self): + ser = SerializeIndependent([['os', True]]) + with pytest.raises(ValueError, match='__call__'): + ser._module_inspect(_NoCall()) + + def test_callable_class_inspect(self): + ser = SerializeIndependent([['os', True]]) + mods = ser._module_inspect(_Adder()) + assert _echo.__module__.split('.')[0] in mods + + def test_partial_inspects_func(self): + ser = SerializeIndependent([['os', True]]) + mods = ser._module_inspect(partial(_echo, 1)) + assert _echo.__module__.split('.')[0] in mods + + def test_dict_iterdata_with_function(self): + ser = SerializeIndependent([['os', True]]) + mods = ser._module_inspect({'fn': _echo, 'n': 1}) + assert _echo.__module__.split('.')[0] in mods + + def test_cython_function_name_inspects_globals(self): + class cython_function_or_method: + __globals__ = {'__file__': '/tmp/foo.pyx'} + + ser = SerializeIndependent([['os', True]]) + mods = ser._module_inspect(cython_function_or_method()) + assert mods == {'/tmp/foo'} + + def test_empty_include_inspects_function(self): + ser = SerializeIndependent([['os', True]]) + strs, paths = ser([_echo], [], ['lithops']) + assert len(strs) == 1 + assert isinstance(paths, set) + + def test_so_origin_added_unless_excluded(self): + ser = SerializeIndependent([['os', True]]) + spec = SimpleNamespace(origin='/opt/ext.so') + with patch('importlib.util.find_spec', return_value=spec): + with patch( + 'lithops.job.serialize.ModuleDependencyAnalyzer' + ) as mda: + mda.return_value.get_and_clear_paths.return_value = set() + _, paths = ser([_echo], [], []) + assert '/opt/ext.so' in paths + + with patch('importlib.util.find_spec', return_value=spec): + with patch( + 'lithops.job.serialize.ModuleDependencyAnalyzer' + ) as mda: + mda.return_value.get_and_clear_paths.return_value = set() + _, excluded = ser([_echo], [], ['ext.so']) + assert '/opt/ext.so' not in excluded + + def test_find_spec_exception_is_swallowed(self): + ser = SerializeIndependent([['os', True]]) + with patch('importlib.util.find_spec', side_effect=ValueError('x')): + with patch( + 'lithops.job.serialize.ModuleDependencyAnalyzer' + ) as mda: + mda.return_value.get_and_clear_paths.return_value = {'/m'} + _, paths = ser([_echo], [], []) + assert '/m' in paths + + +class TestCreateModuleData: + + def test_empty_paths(self): + assert create_module_data(set()) == {} + assert create_module_data([]) == {} + + def test_file_path(self, tmp_path): + f = tmp_path / 'mod.py' + f.write_bytes(b'abc') + data = create_module_data([str(f)]) + assert list(data) == [f.name] + assert data[f.name] == bytes_to_b64str(b'abc') + + def test_directory_collects_nested_py(self, tmp_path): + pkg = tmp_path / 'pkg' + pkg.mkdir() + (pkg / 'a.py').write_bytes(b'a') + nested = pkg / 'sub' + nested.mkdir() + (nested / 'b.py').write_bytes(b'b') + (nested / 'skip.txt').write_bytes(b'no') + data = create_module_data([str(pkg)]) + posix_keys = set(data) + assert posix_keys == {'pkg/a.py', 'pkg/sub/b.py'} + assert all('\\' not in k for k in posix_keys) + + +def _head(headers): + return MagicMock(headers=headers) + + +class _LogCatcher(logging.Handler): + def __init__(self): + super().__init__() + self.messages = [] + + def emit(self, record): + self.messages.append(record.getMessage()) + + +def _catch_partitioner_logs(): + logger = logging.getLogger('lithops.job.partitioner') + handler = _LogCatcher() + logger.addHandler(handler) + prev = logger.level + logger.setLevel(logging.DEBUG) + return logger, handler, prev + + +class TestCreatePartitions: + + def test_empty_iterdata_returns_empty_partitions(self): + assert create_partitions({}, _storage(), [], None, None, '\n') == ( + [], + [], + ) + + def test_http_takes_precedence_over_paths(self, tmp_path): + f = tmp_path / 'f.txt' + f.write_bytes(b'hello world!!') + headers = {'content-length': '13', 'accept-ranges': 'bytes'} + with patch( + 'lithops.job.partitioner.requests.head', return_value=_head(headers) + ): + parts, ppo = create_partitions( + {}, + _storage(), + [{'obj': 'http://example.com/a'}, {'obj': str(f)}], + None, + None, + None, + ) + assert all(isinstance(p['obj'], CloudObjectUrl) for p in parts) + assert ppo == [1] + + def test_https_is_treated_as_url(self): + headers = {'content-length': '10', 'accept-ranges': 'bytes'} + with patch( + 'lithops.job.partitioner.requests.head', return_value=_head(headers) + ): + parts, ppo = create_partitions( + {}, + _storage(), + [{'obj': 'https://example.com/a'}], + None, + None, + None, + ) + assert isinstance(parts[0]['obj'], CloudObjectUrl) + assert ppo == [1] + + def test_url_without_content_length_sets_size_one(self): + with patch( + 'lithops.job.partitioner.requests.head', + return_value=_head({'accept-ranges': 'bytes'}), + ): + parts, ppo = create_partitions( + {}, + _storage(), + [{'obj': 'http://example.com/a'}], + None, + None, + '\n', + ) + assert len(parts) == 1 + assert ppo == [1] + assert parts[0]['obj'].chunk_size == 1 + + def test_url_without_accept_ranges_uses_full_object(self): + with patch( + 'lithops.job.partitioner.requests.head', + return_value=_head({'content-length': '100'}), + ): + parts, ppo = create_partitions( + {}, + _storage(), + [{'obj': 'http://example.com/a'}], + 10, + None, + None, + ) + assert len(parts) == 1 + assert parts[0]['obj'].data_byte_range is None + assert parts[0]['obj'].chunk_size == 100 + assert ppo == [1] + + def test_url_chunk_size_without_newline(self): + headers = {'content-length': '10000', 'accept-ranges': 'bytes'} + with patch( + 'lithops.job.partitioner.requests.head', return_value=_head(headers) + ): + parts, ppo = create_partitions( + {}, + _storage(), + [{'obj': 'http://example.com/a'}], + 4000, + None, + None, + ) + assert ppo == [3] + assert parts[0]['obj'].data_byte_range == (0, 3999) + assert parts[1]['obj'].data_byte_range == (4000, 7999) + assert parts[2]['obj'].data_byte_range == (8000, 11999) + assert parts[0]['obj'].part == 1 + assert parts[0]['obj'].total_parts == 3 + assert parts[0]['obj'].newline is None + + def test_url_chunk_number(self): + headers = {'content-length': '100', 'accept-ranges': 'bytes'} + with patch( + 'lithops.job.partitioner.requests.head', return_value=_head(headers) + ): + parts, ppo = create_partitions( + {}, + _storage(), + [{'obj': 'http://example.com/a'}], + None, + 2, + None, + ) + assert ppo == [2] + assert len(parts) == 2 + + def test_url_newline_uses_chunk_threshold(self): + headers = {'content-length': '10000', 'accept-ranges': 'bytes'} + with patch( + 'lithops.job.partitioner.requests.head', return_value=_head(headers) + ): + parts, _ = create_partitions( + {}, + _storage(), + [{'obj': 'http://example.com/a'}], + 4000, + None, + '\n', + ) + assert parts[0]['obj'].data_byte_range == (0, 4000 + CHUNK_THRESHOLD) + + def test_swapped_chunk_logs(self): + headers = {'content-length': '10', 'accept-ranges': 'bytes'} + logger, handler, prev = _catch_partitioner_logs() + try: + with patch( + 'lithops.job.partitioner.requests.head', + return_value=_head(headers), + ): + create_partitions( + {}, + _storage(), + [{'obj': 'http://example.com/a'}], + None, + 2, + None, + ) + assert 'Chunk number set to 2' in handler.messages + + handler.messages.clear() + with patch( + 'lithops.job.partitioner.requests.head', + return_value=_head(headers), + ): + create_partitions( + {}, + _storage(), + [{'obj': 'http://example.com/a'}], + 10, + None, + None, + ) + assert 'Chunk size set to 10' in handler.messages + + handler.messages.clear() + with patch( + 'lithops.job.partitioner.requests.head', + return_value=_head(headers), + ): + create_partitions( + {}, + _storage(), + [{'obj': 'http://example.com/a'}], + None, + None, + None, + ) + assert 'Chunk size and chunk number not set' in handler.messages + assert 'Chunk size and chunk number not set ' not in handler.messages + finally: + logger.removeHandler(handler) + logger.setLevel(prev) + + def test_paths_file_and_directory(self, tmp_path): + d = tmp_path / 'd' + d.mkdir() + f1 = d / 'a.txt' + f1.write_bytes(b'hello world!!') + nested = d / 'sub' + nested.mkdir() + f2 = tmp_path / 'b.txt' + f2.write_bytes(b'hello world!!') + parts, ppo = create_partitions( + {}, + _storage(), + [{'obj': str(d), 'k': 1}, {'obj': str(f2)}, {'obj': str(f2)}], + None, + None, + None, + ) + paths = {p['obj'].path for p in parts} + assert str(f1) in paths + assert str(f2) in paths + assert str(nested) not in paths + assert all(isinstance(p['obj'], CloudObjectLocal) for p in parts) + for p in parts: + if p['obj'].path == str(f1): + assert p['k'] == 1 + else: + assert 'k' not in p + assert ppo == [1, 1] + + def test_path_one_byte_file_has_one_partition(self, tmp_path): + f = tmp_path / 'tiny.txt' + f.write_bytes(b'x') + parts, ppo = create_partitions( + {}, _storage(), [{'obj': str(f)}], None, None, None + ) + assert len(parts) == 1 + assert ppo == [1] + assert parts[0]['obj'].chunk_size == 1 + + def test_object_storage_head_and_params(self): + internal = _storage() + internal.storage.head_object.return_value = {'content-length': '100'} + parts, ppo = create_partitions( + {}, + internal, + [{'obj': 'localhost://bucket/dir/file.txt', 'n': 7}], + None, + None, + None, + ) + internal.storage.head_object.assert_called_once() + object_key = internal.storage.head_object.call_args[0][1] + assert object_key == 'dir/file.txt' + assert '\\' not in object_key + assert parts[0]['obj'].key == 'dir/file.txt' + assert len(parts) == 1 + assert isinstance(parts[0]['obj'], CloudObject) + assert parts[0]['n'] == 7 + assert parts[0]['obj'].backend == 'localhost' + assert parts[0]['obj'].bucket == 'bucket' + assert ppo == [1] + + def test_object_storage_cloudobject_type_is_converted(self): + internal = _storage() + internal.storage.head_object.return_value = {'content-length': '20'} + co = CloudObject('localhost', 'bucket', 'k') + parts, _ = create_partitions( + {}, internal, [{'obj': co}], None, None, None + ) + assert parts[0]['obj'].key.endswith('k') + + def test_object_storage_missing_scheme_uses_backend(self): + internal = _storage() + internal.storage.list_objects.return_value = [ + {'Key': 'a', 'Size': 20} + ] + parts, _ = create_partitions( + {}, internal, [{'obj': 'bucket'}], None, None, None + ) + internal.storage.list_objects.assert_called_once_with('bucket') + assert parts[0]['obj'].backend == 'localhost' + + def test_object_storage_prefix_listing(self): + internal = _storage() + internal.storage.list_objects.return_value = [ + {'Key': 'dir/a', 'Size': 20} + ] + create_partitions( + {}, + internal, + [{'obj': 'localhost://bucket/dir/'}], + None, + None, + None, + ) + internal.storage.list_objects.assert_called_once() + + def test_object_storage_discard_prefix_folder(self): + internal = _storage() + internal.storage.list_objects.return_value = [ + {'Key': 'dir/', 'Size': 0} + ] + # A listing of nothing but folder markers holds no data at all + with pytest.raises(Exception, match='No objects found'): + create_partitions( + {}, internal, [{'obj': 'localhost://bucket'}], None, None, None + ) + + def test_object_storage_folder_marker_does_not_hide_objects(self): + internal = _storage() + internal.storage.list_objects.return_value = [ + {'Key': 'dir/', 'Size': 0}, + {'Key': 'dir/data.csv', 'Size': 100}, + ] + parts, ppo = create_partitions( + {}, internal, [{'obj': 'localhost://bucket'}], None, None, None + ) + assert len(parts) == 1 + assert ppo == [1] + + def test_object_storage_no_objects_raises(self): + internal = _storage() + internal.storage.list_objects.return_value = [] + with pytest.raises(Exception, match='No objects found'): + create_partitions( + {}, internal, [{'obj': 'localhost://bucket'}], None, None, None + ) + + def test_object_storage_multiple_backends_raises(self): + internal = _storage() + with pytest.raises(Exception, match='multiple storage backends'): + create_partitions( + {}, + internal, + [ + {'obj': 'localhost://b/k'}, + {'obj': 'aws_s3://b/k'}, + ], + None, + None, + None, + ) + + def test_object_storage_other_backend_uses_storage_class(self): + internal = _storage() + fake = MagicMock() + fake.head_object.return_value = {'content-length': '20'} + with patch('lithops.job.partitioner.Storage', return_value=fake) as st: + parts, _ = create_partitions( + {'lithops': {}}, + internal, + [{'obj': 'aws_s3://b/k'}], + None, + None, + None, + ) + st.assert_called_once_with(config={'lithops': {}}, backend='aws_s3') + assert parts[0]['obj'].backend == 'aws_s3' + + def test_object_storage_unset_log_has_no_trailing_space(self): + internal = _storage() + internal.storage.head_object.return_value = {'content-length': '20'} + logger, handler, prev = _catch_partitioner_logs() + try: + create_partitions( + {}, + internal, + [{'obj': 'localhost://bucket/file'}], + None, + None, + None, + ) + assert 'Chunk size and chunk number not set' in handler.messages + assert 'Chunk size and chunk number not set ' not in handler.messages + finally: + logger.removeHandler(handler) + logger.setLevel(prev) + + def test_object_storage_glob_in_obj_name(self): + internal = _storage() + fake = MagicMock() + fake.list_objects.return_value = [{'Key': 'dir/foo1', 'Size': 20}] + with patch('lithops.job.partitioner.Storage', return_value=fake): + parts, _ = create_partitions( + {}, + internal, + [{'obj': 'aws_s3://bucket/dir/foo*'}], + None, + None, + None, + ) + args = fake.list_objects.call_args[0] + assert args[0] == 'bucket' + assert args[1] == 'dir/foo/' + assert args[2] == 'dir/foo*' + assert '\\' not in args[2] + assert len(parts) == 1 + + def test_object_storage_glob_in_prefix(self): + internal = _storage() + fake = MagicMock() + fake.list_objects.return_value = [{'Key': 'pre/a', 'Size': 20}] + with patch('lithops.job.partitioner.Storage', return_value=fake): + create_partitions( + {}, + internal, + [{'obj': 'aws_s3://bucket/pre*/file'}], + None, + None, + None, + ) + args = fake.list_objects.call_args[0] + assert args[0] == 'bucket' + assert args[1] == 'pre/' + assert args[2] == 'pre*/file' + assert '\\' not in args[2] + + def test_object_storage_chunk_number_on_zero_size(self): + internal = _storage() + internal.storage.head_object.return_value = {'content-length': '0'} + parts, ppo = create_partitions( + {}, + internal, + [{'obj': 'localhost://bucket/file'}], + None, + 3, + None, + ) + assert parts == [] + assert ppo == [0] + + +class TestJobExports: + + def test_package_exports(self): + import lithops.job as jobmod + assert jobmod.create_map_job is create_map_job + assert jobmod.create_reduce_job is create_reduce_job + assert jobmod.__all__ == ['create_map_job', 'create_reduce_job'] diff --git a/lithops/tests/test_joblib.py b/lithops/tests/test_joblib.py new file mode 100644 index 000000000..a3ed4a169 --- /dev/null +++ b/lithops/tests/test_joblib.py @@ -0,0 +1,234 @@ +# +# Live tests of the joblib backend, which is what the examples in examples/ +# use to run scikit-learn searches on Lithops. +# +# These run real jobs, so every class skips unless the packages it needs are +# installed, and they pin themselves to the localhost backend: they are here +# to exercise the joblib integration, not to spend money on a cloud one. +# + +import pytest + +joblib = pytest.importorskip('joblib') +pytest.importorskip('diskcache') +pytest.importorskip('numpy') + +from lithops.multiprocessing import config as mp_config # noqa: E402 +from lithops.util.joblib import register_lithops # noqa: E402 + +LOCALHOST_ARGS = {'backend': 'localhost', 'storage': 'localhost'} + + +@pytest.fixture(autouse=True) +def lithops_joblib_backend(): + """ + Registers the backend and gives lithops.multiprocessing its parameters + back afterwards, as they are process-wide + """ + register_lithops() + saved = mp_config.get_parameter(mp_config.LITHOPS_CONFIG) + yield + mp_config.set_parameter(mp_config.LITHOPS_CONFIG, saved) + + +def double(x): + return x * 2 + + +def _on_localhost(**extra): + # n_jobs has to be given: parallel_config leaves it at 1, and joblib runs + # a single job inline without ever asking the backend, unlike the + # parallel_backend the examples use, which defaults it to -1. + # + # A number rather than -1, because -1 asks lithops.multiprocessing for a + # cpu_count, and that reads the default configuration of the machine + # instead of the one these tests pin + extra.setdefault('n_jobs', 4) + return joblib.parallel_config( + backend='lithops', lithops_args=LOCALHOST_ARGS, **extra + ) + + +def _counting_optimizer(collected): + """ + Wraps the shared-object optimizer to record how many calls of each batch + had arguments replaced, which is proof the batch went through Lithops + """ + from lithops.util.joblib import lithops_backend + + real_find = lithops_backend.find_shared_objects + + def counting_find(calls): + out = real_find(calls) + collected.append(sum(1 for call in out if len(call) > 3)) + return out + + return counting_find + + +class TestJoblibBackendLive: + + def test_parallel_over_the_lithops_backend(self): + from unittest.mock import patch + + from lithops.util.joblib import lithops_backend + + batches = [] + with patch.object( + lithops_backend, 'find_shared_objects', + _counting_optimizer(batches) + ): + with _on_localhost(): + results = joblib.Parallel()( + joblib.delayed(double)(i) for i in range(4) + ) + assert results == [0, 2, 4, 6] + # Without this the test would pass on joblib running it inline + assert batches, 'the calls never went through the Lithops backend' + + def test_parallel_with_threads_preferred(self): + # One task runs the whole batch, each call in a thread of its own + with _on_localhost(prefer='threads'): + results = joblib.Parallel()( + joblib.delayed(double)(i) for i in range(3) + ) + assert results == [0, 2, 4] + + def test_an_exception_in_a_call_reaches_the_caller(self): + with _on_localhost(): + with pytest.raises(ZeroDivisionError): + joblib.Parallel()( + joblib.delayed(_divide)(1, d) for d in (1, 0) + ) + + def test_the_shared_argument_travels_once(self): + # What the backend exists for: the same list is an argument of every + # call, so it goes to storage once and the calls carry a reference + import numpy as np + + shared = np.arange(64) + with _on_localhost(): + results = joblib.Parallel()( + joblib.delayed(_sum_with)(shared, i) for i in range(4) + ) + assert results == [int(shared.sum()) + i for i in range(4)] + + def test_lithops_args_pins_the_backend(self): + with _on_localhost(): + assert mp_config.get_parameter(mp_config.LITHOPS_CONFIG) == ( + LOCALHOST_ARGS + ) + + +def _divide(a, b): + return a / b + + +def _sum_with(shared, i): + return int(shared.sum()) + i + + +class TestSklearnOverJoblib: + """ + The searches the examples in examples/ run, in a smaller shape so that + they finish in seconds on the localhost backend + """ + + @pytest.fixture(autouse=True) + def _needs_sklearn(self): + pytest.importorskip('sklearn') + + def test_grid_search_over_the_lithops_backend(self): + from sklearn.datasets import load_digits + from sklearn.model_selection import GridSearchCV + from sklearn.tree import DecisionTreeClassifier + + digits = load_digits() + search = GridSearchCV( + DecisionTreeClassifier(random_state=0), + {'max_depth': [2, 4]}, + cv=2, + refit=True, + ) + + with _on_localhost(): + search.fit(digits.data, digits.target) + + assert search.best_params_['max_depth'] in (2, 4) + assert 0.0 < search.best_score_ <= 1.0 + # refit ran, so the search can predict + assert len(search.predict(digits.data[:5])) == 5 + + def test_the_dataset_is_proxied_for_every_fit_of_the_search(self): + # Every fit gets the same X and y, so they travel as one cloud object + # instead of once per task. A call whose arguments were replaced + # carries a fourth element with their positions + from unittest.mock import patch + + from sklearn.datasets import load_digits + from sklearn.model_selection import GridSearchCV + from sklearn.tree import DecisionTreeClassifier + + from lithops.util.joblib import lithops_backend + + digits = load_digits() + search = GridSearchCV( + DecisionTreeClassifier(random_state=0), + {'max_depth': [2, 4, 6]}, + cv=2, + ) + + proxied = [] + with patch.object( + lithops_backend, 'find_shared_objects', + _counting_optimizer(proxied) + ): + with _on_localhost(): + search.fit(digits.data, digits.target) + + assert proxied, 'the batch never went through the optimizer' + # Three candidates over two folds + assert max(proxied) >= 6 + + def test_randomized_search_over_the_lithops_backend(self): + import numpy as np + from sklearn.datasets import load_digits + from sklearn.model_selection import RandomizedSearchCV + from sklearn.tree import DecisionTreeClassifier + + digits = load_digits() + search = RandomizedSearchCV( + DecisionTreeClassifier(random_state=0), + {'min_samples_leaf': np.arange(1, 10)}, + cv=2, + n_iter=3, + random_state=0, + ) + + with _on_localhost(): + search.fit(digits.data, digits.target) + + assert 0.0 < search.best_score_ <= 1.0 + + def test_a_pipeline_search_over_the_lithops_backend(self): + # The shape of examples/sklearn_job_3.py, without pandas + from sklearn.datasets import load_digits + from sklearn.model_selection import GridSearchCV + from sklearn.pipeline import Pipeline + from sklearn.preprocessing import StandardScaler + from sklearn.tree import DecisionTreeClassifier + + digits = load_digits() + pipeline = Pipeline([ + ('scale', StandardScaler()), + ('classifier', DecisionTreeClassifier(random_state=0)), + ]) + search = GridSearchCV( + pipeline, {'classifier__max_depth': [2, 4]}, cv=2, refit=True + ) + + with _on_localhost(): + search.fit(digits.data, digits.target) + + assert search.best_params_['classifier__max_depth'] in (2, 4) + assert 0.0 < search.best_score_ <= 1.0 diff --git a/lithops/tests/test_localhost.py b/lithops/tests/test_localhost.py new file mode 100644 index 000000000..14378757e --- /dev/null +++ b/lithops/tests/test_localhost.py @@ -0,0 +1,1002 @@ +# +# Unit tests for the localhost compute and storage backends. +# No Docker daemon required; subprocess and docker CLI calls are mocked. +# + +import copy +import io +import json +import logging +import os +import queue +import shutil +import signal +import subprocess as sp +import sys +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +import lithops +from lithops.constants import JOBS_PREFIX, TEMP_DIR, USER_TEMP_DIR +from lithops.localhost import LocalhostHandler, LocalhostHandlerV1, LocalhostHandlerV2 +from lithops.localhost import config as localhost_config +from lithops.localhost.v1 import localhost as v1 +from lithops.localhost.v1 import runner as v1_runner +from lithops.localhost.v2 import localhost as v2 +from lithops.localhost.v2 import runner as v2_runner +from lithops.localhost.utils import ( + copy_lithops_package, + decode_process_output, + docker_pull_cmd, + docker_rm_cmd, + docker_run_cmd, + log_process_failure, +) +from lithops.storage.backends.localhost.localhost import LocalhostStorageBackend +from lithops.storage.utils import StorageNoSuchKeyError +from lithops.tests.functions import simple_map_function, sleep_seconds +from lithops.utils import BackendType, CountDownLatch +from lithops.version import __version__ + + +def _config(**extra): + cfg = { + 'runtime': 'python3', + 'max_workers': 1, + 'worker_processes': 2, + } + cfg.update(extra) + return cfg + + +def _job_payload(**extra): + payload = { + 'executor_id': 'sess-0', + 'job_id': 'M000', + 'job_key': 'sess-0-M000', + 'call_ids': ['00000', '00001'], + 'data_byte_ranges': [(0, 1), (2, 3)], + 'config': { + 'lithops': {'storage': 'localhost'}, + 'localhost': {'storage_bucket': 'storage'}, + }, + } + payload.update(extra) + return payload + + +class TestLocalhostConfig: + + def test_environment_default_for_python_and_paths(self): + for runtime in ( + 'python', + 'python3', + 'python3.12', + '/usr/bin/python3', + r'C:\Python\python.exe', + ): + assert localhost_config.get_environment(runtime) is ( + localhost_config.LocalhostEnvironment.DEFAULT + ) + + def test_environment_container_for_image_names(self): + for runtime in ( + 'lithopscloud/ibmcf-python-v312', + 'docker.io/lithopscloud/ibmcf-python-v312', + ): + assert localhost_config.get_environment(runtime) is ( + localhost_config.LocalhostEnvironment.CONTAINER + ) + + def test_environment_python_tagged_image_is_container(self): + assert localhost_config.get_environment('python:3.12') is ( + localhost_config.LocalhostEnvironment.CONTAINER + ) + + def test_environment_enum_values(self): + assert localhost_config.LocalhostEnvironment.DEFAULT.value == 'default' + assert localhost_config.LocalhostEnvironment.CONTAINER.value == 'container' + + def test_runtime_key_and_info(self): + assert localhost_config.runtime_key('/python3/') == ( + f'localhost/{__version__}/python3' + ) + assert localhost_config.runtime_key(r'C:\Python\python.exe') == ( + f'localhost/{__version__}/C:/Python/python.exe' + ) + assert '\\' not in localhost_config.runtime_key(r'C:\Python\python.exe') + assert localhost_config.runtime_info(_config(runtime_memory=256)) == { + 'runtime_name': 'python3', + 'runtime_memory': 256, + 'runtime_timeout': None, + 'max_workers': 1, + } + + def test_load_config_fills_defaults_and_forces_max_workers(self): + cfg = {'lithops': {}, 'localhost': {'max_workers': 8, 'runtime': 'python3.11'}} + localhost_config.load_config(cfg) + assert cfg['localhost']['runtime'] == 'python3.11' + assert cfg['localhost']['worker_processes'] == (os.cpu_count() or 1) + assert cfg['localhost']['max_workers'] == 1 + assert cfg['lithops']['execution_timeout'] == ( + localhost_config.LOCALHOST_EXECUTION_TIMEOUT + ) + assert cfg['lithops']['storage'] == 'localhost' + + def test_load_config_creates_localhost_section(self): + cfg = {'lithops': {'storage': 's3', 'execution_timeout': 9}} + localhost_config.load_config(cfg) + assert cfg['localhost']['runtime'] == localhost_config.DEFAULT_CONFIG_KEYS['runtime'] + assert cfg['lithops']['storage'] == 's3' + assert cfg['lithops']['execution_timeout'] == 9 + + def test_default_handler_is_v2(self): + assert LocalhostHandler is LocalhostHandlerV2 + + +class TestLocalhostHandlerV2: + + def test_backend_type_runtime_key_and_info(self): + handler = LocalhostHandlerV2(_config(runtime_memory=256, runtime_timeout=60)) + assert handler.get_backend_type() == BackendType.BATCH.value + assert handler.get_runtime_key('/python3/') == ( + f'localhost/{__version__}/python3' + ) + assert handler.get_runtime_info() == { + 'runtime_name': 'python3', + 'runtime_memory': 256, + 'runtime_timeout': 60, + 'max_workers': 1, + } + handler.clean() + + def test_init_selects_default_environment(self): + handler = LocalhostHandlerV2(_config()) + with patch.object(v2, 'DefaultEnvironment') as env_cls: + handler.init() + env_cls.assert_called_once_with(handler.config) + env_cls.return_value.setup.assert_called_once() + assert handler.env is env_cls.return_value + + def test_init_selects_container_environment(self): + handler = LocalhostHandlerV2(_config(runtime='lithops/python:3.12')) + with patch.object(v2, 'ContainerEnvironment') as env_cls: + handler.init() + env_cls.assert_called_once_with(handler.config) + assert handler.env is env_cls.return_value + + def test_invoke_runs_job_and_starts_manager(self): + handler = LocalhostHandlerV2(_config()) + handler.env = MagicMock() + handler.start_manager = MagicMock() + payload = _job_payload() + handler.invoke(payload) + handler.env.run_job.assert_called_once_with(payload) + handler.start_manager.assert_called_once() + assert handler.invocation_in_progress is False + + def test_start_manager_is_noop_when_already_running(self): + handler = LocalhostHandlerV2(_config()) + handler.env = MagicMock() + handler.job_manager = object() + handler.start_manager() + handler.env.start.assert_not_called() + + def test_clear_drains_queue_and_unlocks_jobs(self): + handler = LocalhostHandlerV2(_config()) + handler.env = MagicMock() + handler.env.work_queue = queue.Queue() + handler.env.work_queue.put('task-a') + handler.env.work_queue.put('task-b') + latch = CountDownLatch(2) + handler.env.jobs = {'sess-0-M000': latch} + handler.clear() + handler.env.drop_pending_tasks.assert_called_once_with(None) + handler.env.stop.assert_called_once_with(None) + assert latch.done is True + + def test_clear_leaves_the_latches_of_other_jobs_alone(self): + handler = LocalhostHandlerV2(_config()) + handler.env = MagicMock() + mine, theirs = CountDownLatch(1), CountDownLatch(1) + handler.env.jobs = {'sess-0-M000': mine, 'sess-0-M001': theirs} + handler.clear({'sess-0-M000'}) + assert mine.done is True + assert theirs.done is False + + +class TestLocalhostHandlerV1: + + def test_backend_type_and_runtime_key_match_v2(self): + handler = LocalhostHandlerV1(_config()) + assert handler.get_backend_type() == BackendType.BATCH.value + assert handler.get_runtime_key('python3') == ( + f'localhost/{__version__}/python3' + ) + + def test_init_selects_default_environment(self): + handler = LocalhostHandlerV1(_config()) + with patch.object(v1, 'DefaultEnvironment') as env_cls: + handler.init() + env_cls.assert_called_once_with(handler.config) + env_cls.return_value.setup.assert_called_once() + + def test_invoke_queues_prepared_job_file(self): + handler = LocalhostHandlerV1(_config()) + handler.env = MagicMock() + handler.env.prepare_job_file.return_value = '/tmp/job.json' + handler.start_manager = MagicMock() + payload = _job_payload() + handler.invoke(payload) + handler.env.prepare_job_file.assert_called_once_with(payload) + assert handler.job_queue.get_nowait() == (payload, '/tmp/job.json') + handler.start_manager.assert_called_once() + assert handler.invocation_in_progress is False + + def test_clear_drains_queue_and_sends_sentinel(self): + handler = LocalhostHandlerV1(_config()) + handler.env = MagicMock() + handler.job_manager = object() + handler.job_queue.put(('job', 'file')) + handler.clear() + handler.env.stop.assert_called_once() + assert handler.job_queue.get_nowait() == (None, None) + + def test_clear_leaves_the_queued_jobs_of_others_alone(self): + handler = LocalhostHandlerV1(_config()) + handler.env = MagicMock() + handler.job_manager = object() + mine = ({'job_key': 'sess-0-M000'}, 'mine.json') + theirs = ({'job_key': 'sess-0-M001'}, 'theirs.json') + handler.job_queue.put(mine) + handler.job_queue.put(theirs) + handler.clear({'sess-0-M000'}) + assert handler.job_queue.get_nowait() == theirs + assert handler.job_queue.get_nowait() == (None, None) + assert handler.job_queue.empty() + + +class TestV2Environment: + + def test_run_job_splits_payload_per_call(self, tmp_path, monkeypatch): + monkeypatch.setattr('lithops.localhost.v2.localhost.JOBS_DIR', str(tmp_path)) + env = v2.ExecutionEnvironment(_config(worker_processes=1)) + env.run_job(_job_payload()) + assert env.jobs['sess-0-M000'].count == 2 + first = json.loads(env.work_queue.get_nowait()) + second = json.loads(env.work_queue.get_nowait()) + assert first['call_ids'] == ['00000'] + assert first['data_byte_ranges'] == [[0, 1]] + assert second['call_ids'] == ['00001'] + assert second['data_byte_ranges'] == [[2, 3]] + + def test_start_is_noop_when_consumers_already_running(self): + env = v2.ExecutionEnvironment(_config(worker_processes=1)) + existing = object() + env.consumer_threads = [existing] + env.start() + assert env.consumer_threads == [existing] + + def test_consumer_writes_task_file_then_unlocks(self, tmp_path, monkeypatch): + monkeypatch.setattr('lithops.localhost.v2.localhost.JOBS_DIR', str(tmp_path)) + env = v2.ExecutionEnvironment(_config(worker_processes=1)) + env.run_task = MagicMock() + env.run_job(_job_payload(call_ids=['00000'], data_byte_ranges=[(0, 1)])) + env.start() + try: + env.jobs['sess-0-M000'].wait() + env.run_task.assert_called_once_with('sess-0-M000', '00000') + assert not (tmp_path / 'sess-0-M000' / '00000.task').exists() + finally: + env.stop() + assert env.consumer_threads == [] + + def test_copy_lithops_skips_inside_worker_when_runner_exists( + self, tmp_path, monkeypatch + ): + runner = tmp_path / 'localhost-runner.py' + runner.write_text('# runner\n') + monkeypatch.setattr('lithops.localhost.v2.localhost.RUNNER_FILE', str(runner)) + monkeypatch.setattr( + 'lithops.localhost.v2.localhost.is_lithops_worker', lambda: True + ) + env = v2.ExecutionEnvironment(_config()) + with patch('lithops.localhost.v2.localhost.copy_lithops_package') as copy_pkg: + env._copy_lithops_to_tmp() + copy_pkg.assert_not_called() + + def test_copy_lithops_uses_v2_runner(self, tmp_path, monkeypatch): + monkeypatch.setattr( + 'lithops.localhost.v2.localhost.LITHOPS_TEMP_DIR', str(tmp_path) + ) + monkeypatch.setattr( + 'lithops.localhost.v2.localhost.RUNNER_FILE', + str(tmp_path / 'localhost-runner.py'), + ) + monkeypatch.setattr( + 'lithops.localhost.v2.localhost.is_lithops_worker', lambda: False + ) + env = v2.ExecutionEnvironment(_config()) + with patch('lithops.localhost.v2.localhost.copy_lithops_package') as copy_pkg: + env._copy_lithops_to_tmp() + runner_src = copy_pkg.call_args[0][1] + assert runner_src.endswith(os.path.join('localhost', 'v2', 'runner.py')) + assert copy_pkg.call_args[0][2] == str(tmp_path / 'localhost-runner.py') + assert copy_pkg.call_args[0][3] == str(tmp_path) + + def test_default_run_task_invokes_runner_and_drops_process(self): + env = v2.DefaultEnvironment(_config()) + proc = MagicMock() + proc.returncode = 1 + proc.communicate.return_value = (b'', b'Traceback: boom\n') + with patch('lithops.localhost.v2.localhost.sp.Popen', return_value=proc) as popen, \ + patch('lithops.localhost.v2.localhost.log_process_failure') as log_fail: + env.run_task('sess-0-M000', '00000') + cmd = popen.call_args[0][0] + assert cmd[0] == 'python3' + assert cmd[2:] == [ + 'run_job', + os.path.join(v2.JOBS_DIR, 'sess-0-M000', '00000.task'), + ] + proc.communicate.assert_called_once() + assert 'sess-0-M000-00000' not in env.task_processes + log_fail.assert_called_once() + assert log_fail.call_args.kwargs['stderr'] == b'Traceback: boom\n' + + def test_default_stop_kills_matching_process_group(self): + env = v2.DefaultEnvironment(_config(worker_processes=1)) + proc = MagicMock() + proc.poll.return_value = None + proc.pid = 123 + env.task_processes['sess-0-M000-00000'] = proc + env.jobs = {'sess-0-M000': CountDownLatch(0)} + env.is_unix_system = True + with patch('lithops.localhost.utils.os.getpgid', return_value=9), \ + patch('lithops.localhost.utils.os.killpg') as killpg, \ + patch.object(v2.ExecutionEnvironment, '_teardown'): + env.stop(['sess-0-M000']) + killpg.assert_called_once_with(9, signal.SIGKILL) + assert 'sess-0-M000-00000' not in env.task_processes + + def test_teardown_sends_one_sentinel_per_running_consumer(self): + env = v2.DefaultEnvironment(_config(worker_processes=3)) + # No consumer running: a sentinel would sit in the queue and kill the + # next consumer that starts + env.stop() + assert env.work_queue.empty() + + env.consumer_threads = [MagicMock(), MagicMock()] + env._teardown() + assert env.work_queue.qsize() == 2 + assert env.consumer_threads == [] + + def test_stop_keeps_consumers_while_another_job_runs(self): + env = v2.DefaultEnvironment(_config(worker_processes=1)) + env.jobs = { + 'sess-0-M000': CountDownLatch(0), + 'sess-0-M001': CountDownLatch(1), + } + threads = [MagicMock()] + env.consumer_threads = list(threads) + env.stop(['sess-0-M000']) + assert env.consumer_threads == threads + assert env.work_queue.empty() + threads[0].join.assert_not_called() + + # Once that job is done too, the environment is torn down + env.jobs['sess-0-M001'].unlock() + env.stop(['sess-0-M001']) + assert env.consumer_threads == [] + + def test_drop_pending_tasks_keeps_the_other_jobs_tasks(self): + env = v2.DefaultEnvironment(_config(worker_processes=1)) + mine = json.dumps({'job_key': 'sess-0-M000'}) + theirs = json.dumps({'job_key': 'sess-0-M001'}) + env.work_queue.put(mine) + env.work_queue.put(theirs) + env.work_queue.put(None) + env.drop_pending_tasks({'sess-0-M000'}) + assert env.work_queue.get_nowait() == theirs + assert env.work_queue.empty() + + def test_drop_pending_tasks_empties_the_queue_for_every_job(self): + env = v2.DefaultEnvironment(_config(worker_processes=1)) + env.work_queue.put(json.dumps({'job_key': 'sess-0-M000'})) + env.work_queue.put(None) + env.drop_pending_tasks() + assert env.work_queue.empty() + + def test_task_process_not_started_when_job_was_stopped(self): + env = v2.DefaultEnvironment(_config(worker_processes=1)) + env.stopped_jobs.add('sess-0-M000') + with patch('lithops.localhost.v2.localhost.sp.Popen') as popen: + env._run_task_process('sess-0-M000-00000', ['cmd']) + popen.assert_not_called() + assert 'sess-0-M000-00000' not in env.task_processes + + def test_stop_marks_the_job_and_run_job_clears_it(self): + env = v2.DefaultEnvironment(_config(worker_processes=1)) + env.jobs = {'sess-0-M000': CountDownLatch(1)} + with patch.object(v2.ExecutionEnvironment, '_teardown'): + env.stop(['sess-0-M000']) + assert 'sess-0-M000' in env.stopped_jobs + + payload = { + 'job_key': 'sess-0-M000', + 'call_ids': ['00000'], + 'data_byte_ranges': [None], + } + with patch('lithops.localhost.v2.localhost.os.makedirs'): + env.run_job(payload) + assert 'sess-0-M000' not in env.stopped_jobs + + def test_container_metadata_command_uses_runner_and_user(self): + with patch.object(v2, 'get_docker_path', return_value='/bin/docker'), \ + patch.object(v2, 'is_podman', return_value=False): + env = v2.ContainerEnvironment(_config(runtime='img:tag')) + env.is_unix_system = True + env.uid = 1000 + env.gid = 1000 + result = MagicMock() + result.stdout = '{"preinstalls": []}\n' + with patch('lithops.localhost.v2.localhost.os.path.isfile', return_value=True), \ + patch('lithops.localhost.v2.localhost.sp.run', return_value=result) as run: + assert env.get_metadata() == {'preinstalls': []} + cmd = run.call_args[0][0] + joined = ' '.join(cmd) + assert cmd[0] == '/bin/docker' + assert '--user' in cmd + assert '1000:1000' in joined + assert 'get_metadata' in joined + assert f'/tmp/{USER_TEMP_DIR}/localhost-runner.py' in joined + + def test_container_run_task_uses_docker_exec(self): + with patch.object(v2, 'get_docker_path', return_value='docker'), \ + patch.object(v2, 'is_podman', return_value=False): + env = v2.ContainerEnvironment(_config(runtime='img:tag')) + proc = MagicMock() + proc.returncode = 0 + proc.communicate.return_value = (b'', b'') + with patch('lithops.localhost.v2.localhost.sp.Popen', return_value=proc) as popen: + env.run_task('sess-0-M000', '00000') + joined = ' '.join(popen.call_args[0][0]) + assert f'docker exec {env.container_name}' in joined + assert 'run_job' in joined + assert f'/tmp/{USER_TEMP_DIR}/jobs/sess-0-M000/00000.task' in joined + assert 'sess-0-M000-00000' not in env.task_processes + + def test_container_setup_pulls_with_docker_path(self): + with patch.object(v2, 'get_docker_path', return_value='/bin/podman'), \ + patch.object(v2, 'is_podman', return_value=True): + env = v2.ContainerEnvironment( + _config(runtime='img:tag', pull_runtime=True) + ) + with patch.object(env, '_copy_lithops_to_tmp'), \ + patch('lithops.localhost.v2.localhost.sp.run') as run: + env.setup() + assert run.call_args[0][0][:3] == ['/bin/podman', 'pull', 'img:tag'] + + def test_container_docker_volume_is_posix_and_skips_user_on_windows(self): + with patch.object(v2, 'get_docker_path', return_value='docker'), \ + patch.object(v2, 'is_podman', return_value=False), \ + patch.object(v2, 'is_unix_system', return_value=False): + env = v2.ContainerEnvironment(_config(runtime='img:tag')) + assert env.uid is None + assert env.gid is None + cmd = env._container_run_cmd('lithops_win') + assert '--user' not in cmd + volume = cmd[cmd.index('-v') + 1] + assert volume == f'{Path(TEMP_DIR).as_posix()}:/tmp' + assert '\\' not in volume + + +class TestV1Environment: + + def test_job_process_not_started_when_job_was_stopped(self): + env = v1.DefaultEnvironment(_config()) + env.stopped_jobs.add('sess-0-M000') + with patch('lithops.localhost.v1.localhost.sp.Popen') as popen: + assert env._start_job_process('sess-0-M000', ['cmd']) is None + popen.assert_not_called() + assert 'sess-0-M000' not in env.jobs + + def test_stop_marks_the_job_and_prepare_job_file_clears_it( + self, tmp_path, monkeypatch + ): + monkeypatch.setattr( + 'lithops.localhost.v1.localhost.LITHOPS_TEMP_DIR', str(tmp_path) + ) + env = v1.DefaultEnvironment(_config()) + proc = MagicMock() + proc.pid = 42 + env.jobs = {'sess-0-M000': proc} + with patch('lithops.localhost.utils.kill_process'): + env.stop(['sess-0-M000']) + assert 'sess-0-M000' in env.stopped_jobs + assert 'sess-0-M000' not in env.jobs + + env.prepare_job_file(_job_payload()) + assert 'sess-0-M000' not in env.stopped_jobs + + def test_prepare_job_file_writes_payload_and_returns_local_path( + self, tmp_path, monkeypatch + ): + monkeypatch.setattr( + 'lithops.localhost.v1.localhost.LITHOPS_TEMP_DIR', str(tmp_path) + ) + env = v1.DefaultEnvironment(_config()) + payload = _job_payload() + filename = env.prepare_job_file(payload) + assert filename == os.path.join( + str(tmp_path), 'storage', JOBS_PREFIX, 'sess-0-M000-job.json' + ) + with open(filename) as fh: + assert json.load(fh)['job_key'] == 'sess-0-M000' + + def test_prepare_job_file_returns_docker_path_for_container( + self, tmp_path, monkeypatch + ): + monkeypatch.setattr( + 'lithops.localhost.v1.localhost.LITHOPS_TEMP_DIR', str(tmp_path) + ) + with patch.object(v1, 'get_docker_path', return_value='docker'), \ + patch.object(v1, 'is_podman', return_value=False): + env = v1.ContainerEnvironment(_config(runtime='img:tag')) + filename = env.prepare_job_file(_job_payload()) + assert filename == ( + f'/tmp/{USER_TEMP_DIR}/storage/{JOBS_PREFIX}/sess-0-M000-job.json' + ) + assert '\\' not in filename + local = os.path.join( + str(tmp_path), 'storage', JOBS_PREFIX, 'sess-0-M000-job.json' + ) + assert os.path.isfile(local) + + def test_default_run_job_stores_process(self): + env = v1.DefaultEnvironment(_config()) + proc = MagicMock() + with patch('lithops.localhost.v1.localhost.os.path.isfile', return_value=True), \ + patch('lithops.localhost.v1.localhost.sp.Popen', return_value=proc) as popen: + assert env.run_job('sess-0-M000', '/tmp/job.json') is proc + assert popen.call_args[0][0] == [ + 'python3', v1.RUNNER_FILE, 'run_job', '/tmp/job.json' + ] + assert env.jobs['sess-0-M000'] is proc + + def test_container_run_job_names_container_after_job_key(self): + with patch.object(v1, 'get_docker_path', return_value='docker'), \ + patch.object(v1, 'is_podman', return_value=False): + env = v1.ContainerEnvironment( + _config(runtime='img:tag', use_gpu=True) + ) + proc = MagicMock() + with patch('lithops.localhost.v1.localhost.os.path.isfile', return_value=True), \ + patch('lithops.localhost.v1.localhost.sp.Popen', return_value=proc) as popen: + env.run_job('sess-0-M000', '/tmp/job.json') + joined = ' '.join(popen.call_args[0][0]) + assert '--name lithops_sess-0-M000' in joined + assert '--gpus all' in joined + assert 'run_job /tmp/job.json' in joined + assert env.jobs['sess-0-M000'] is proc + + def test_container_setup_pulls_with_docker_path(self): + with patch.object(v1, 'get_docker_path', return_value='/bin/podman'), \ + patch.object(v1, 'is_podman', return_value=True): + env = v1.ContainerEnvironment( + _config(runtime='img:tag', pull_runtime=True) + ) + with patch.object(env, '_copy_lithops_to_tmp'), \ + patch('lithops.localhost.v1.localhost.sp.run') as run: + env.setup() + assert run.call_args[0][0][:3] == ['/bin/podman', 'pull', 'img:tag'] + + def test_container_docker_volume_is_posix_and_skips_user_on_windows(self): + with patch.object(v1, 'get_docker_path', return_value='docker'), \ + patch.object(v1, 'is_podman', return_value=False), \ + patch.object(v1, 'is_unix_system', return_value=False): + env = v1.ContainerEnvironment(_config(runtime='img:tag')) + assert env.uid is None + assert env.gid is None + cmd = env._container_cmd('lithops_win', ['/tmp/job.json']) + assert '--user' not in cmd + volume = cmd[cmd.index('-v') + 1] + assert volume == f'{Path(TEMP_DIR).as_posix()}:/tmp' + assert '\\' not in volume + assert '/tmp/job.json' in cmd + + def test_copy_lithops_uses_v1_runner(self, tmp_path, monkeypatch): + monkeypatch.setattr( + 'lithops.localhost.v1.localhost.LITHOPS_TEMP_DIR', str(tmp_path) + ) + monkeypatch.setattr( + 'lithops.localhost.v1.localhost.RUNNER_FILE', + str(tmp_path / 'localhost-runner.py'), + ) + monkeypatch.setattr( + 'lithops.localhost.v1.localhost.is_lithops_worker', lambda: False + ) + env = v1.ExecutionEnvironment(_config()) + with patch('lithops.localhost.v1.localhost.copy_lithops_package') as copy_pkg: + env._copy_lithops_to_tmp() + runner_src = copy_pkg.call_args[0][1] + assert runner_src.endswith(os.path.join('localhost', 'v1', 'runner.py')) + + def test_stop_kills_process_group(self): + env = v1.DefaultEnvironment(_config()) + proc = MagicMock() + proc.poll.return_value = None + proc.pid = 77 + env.jobs['sess-0-M000'] = proc + env.is_unix_system = True + with patch('lithops.localhost.utils.os.getpgid', return_value=5), \ + patch('lithops.localhost.utils.os.killpg') as killpg: + env.stop(['sess-0-M000']) + killpg.assert_called_once_with(5, signal.SIGKILL) + assert 'sess-0-M000' not in env.jobs + + +class TestLocalhostUtils: + + def test_docker_command_helpers(self): + assert docker_pull_cmd('/bin/podman', 'img:tag') == [ + '/bin/podman', 'pull', 'img:tag' + ] + assert docker_rm_cmd('docker', 'lithops_abc') == [ + 'docker', 'rm', '-f', 'lithops_abc' + ] + + def test_docker_run_cmd_volume_is_posix_and_user_is_unix_only(self): + cmd = docker_run_cmd( + 'docker', + 'img:tag', + name='lithops_job', + tmp_path=Path(TEMP_DIR).as_posix(), + ) + volume = cmd[cmd.index('-v') + 1] + assert volume == f'{Path(TEMP_DIR).as_posix()}:/tmp' + assert '\\' not in volume + assert '--user' not in cmd + + unix_cmd = docker_run_cmd( + 'docker', + 'img:tag', + name='lithops_job', + tmp_path='/var/folders/xx/T', + uid=1000, + gid=1000, + ) + assert unix_cmd[unix_cmd.index('--user') + 1] == '1000:1000' + assert '-v' in unix_cmd + assert unix_cmd[unix_cmd.index('-v') + 1] == '/var/folders/xx/T:/tmp' + + def test_copy_lithops_package_skips_pycache(self, tmp_path): + src = tmp_path / 'src' / 'lithops' + src.mkdir(parents=True) + (src / 'mod.py').write_text('x = 1\n') + cache = src / '__pycache__' + cache.mkdir() + (cache / 'mod.cpython-312.pyc').write_bytes(b'pyc') + runner_src = tmp_path / 'runner.py' + runner_src.write_text('# runner\n') + dest_root = tmp_path / 'dest' + copy_lithops_package( + str(src), str(runner_src), str(dest_root / 'runner.py'), str(dest_root) + ) + copied = dest_root / 'lithops' + assert (copied / 'mod.py').is_file() + assert not (copied / '__pycache__').exists() + assert (dest_root / 'runner.py').read_text() == '# runner\n' + + +class TestLocalhostRunners: + + def test_import_does_not_open_log_stream(self): + assert getattr(v1_runner, 'log_file_stream', None) is None + assert getattr(v2_runner, 'log_file_stream', None) is None + + def _patch_runner_paths(self, runner, tmp_path, monkeypatch): + monkeypatch.setattr(runner, 'LITHOPS_TEMP_DIR', str(tmp_path)) + monkeypatch.setattr(runner, 'JOBS_DIR', str(tmp_path / 'jobs')) + monkeypatch.setattr(runner, 'LOGS_DIR', str(tmp_path / 'logs')) + monkeypatch.setattr(runner, 'RN_LOG_FILE', str(tmp_path / 'runner.log')) + monkeypatch.setattr(runner, '_set_fork_start_method', lambda: None) + + def test_v1_unknown_command_exits(self, tmp_path, monkeypatch): + self._patch_runner_paths(v1_runner, tmp_path, monkeypatch) + monkeypatch.setattr(sys, 'argv', ['runner.py', 'not-a-command']) + with pytest.raises(SystemExit) as exc: + v1_runner.main() + assert exc.value.code == 1 + + def test_v2_unknown_command_exits(self, tmp_path, monkeypatch): + self._patch_runner_paths(v2_runner, tmp_path, monkeypatch) + monkeypatch.setattr(sys, 'argv', ['runner.py', 'not-a-command']) + with pytest.raises(SystemExit) as exc: + v2_runner.main() + assert exc.value.code == 1 + + def test_v1_run_job_reads_text_json(self, tmp_path, monkeypatch): + self._patch_runner_paths(v1_runner, tmp_path, monkeypatch) + job_file = tmp_path / 'job.json' + job_file.write_text(json.dumps(_job_payload(call_ids=['00000']))) + monkeypatch.setattr(sys, 'argv', ['runner.py', 'run_job', str(job_file)]) + with patch.object(v1_runner, 'function_handler') as handler: + v1_runner.main() + handler.assert_called_once() + assert handler.call_args[0][0]['job_key'] == 'sess-0-M000' + assert not job_file.exists() + assert (tmp_path / 'jobs' / 'sess-0-M000.done').is_file() + + def test_v2_run_job_forces_single_worker_process(self, tmp_path, monkeypatch): + self._patch_runner_paths(v2_runner, tmp_path, monkeypatch) + task_file = tmp_path / '00000.task' + task_file.write_text(json.dumps(_job_payload(call_ids=['00000']))) + monkeypatch.setattr(sys, 'argv', ['runner.py', 'run_job', str(task_file)]) + with patch.object(v2_runner, 'function_handler') as handler: + v2_runner.main() + assert handler.call_args[0][0]['worker_processes'] == 1 + + def test_v1_run_job_prints_exception_on_failure(self, tmp_path, monkeypatch): + self._patch_runner_paths(v1_runner, tmp_path, monkeypatch) + job_file = tmp_path / 'job.json' + job_file.write_text(json.dumps(_job_payload(call_ids=['00000']))) + monkeypatch.setattr(sys, 'argv', ['runner.py', 'run_job', str(job_file)]) + err = io.StringIO() + monkeypatch.setattr(v1_runner.sys, '__stderr__', err) + with patch.object( + v1_runner, 'function_handler', side_effect=RuntimeError('boom') + ): + with pytest.raises(SystemExit) as exc: + v1_runner.main() + assert exc.value.code == 1 + assert 'RuntimeError: boom' in err.getvalue() + + +class TestLogProcessFailure: + + def test_decode_bytes_and_str(self): + assert decode_process_output(b' err \n') == 'err' + assert decode_process_output(' out ') == 'out' + assert decode_process_output(None) == '' + assert decode_process_output(object()) == '' + + def test_logs_stderr_before_runner_file(self, caplog, tmp_path): + log_file = tmp_path / 'runner.log' + log_file.write_text('stale runner output\n') + with caplog.at_level(logging.ERROR): + log_process_failure( + logging.getLogger('test-fail'), + 'process failed with return code 1', + stdout=b'ignored stdout', + stderr=b'Traceback: boom', + log_file=str(log_file), + ) + assert 'process failed with return code 1' in caplog.text + assert 'Traceback: boom' in caplog.text + assert 'stale runner output' not in caplog.text + + def test_falls_back_to_runner_log_tail(self, caplog, tmp_path): + log_file = tmp_path / 'runner.log' + log_file.write_text('ModuleNotFoundError: numcodecs\n') + with caplog.at_level(logging.ERROR): + log_process_failure( + logging.getLogger('test-fail'), + 'process failed with return code 1', + stdout=b'', + stderr=b'', + log_file=str(log_file), + ) + assert 'Runner log' in caplog.text + assert 'numcodecs' in caplog.text + + +class TestLocalhostStorageBackend: + + @pytest.fixture + def backend(self, tmp_path, monkeypatch): + monkeypatch.setattr( + 'lithops.storage.backends.localhost.localhost.LITHOPS_TEMP_DIR', + str(tmp_path), + ) + return LocalhostStorageBackend({}), tmp_path + + def test_put_get_bytes_str_and_stream(self, backend): + storage, _ = backend + storage.put_object('bucket', 'bytes.bin', b'abc') + storage.put_object('bucket', 'text.txt', 'hello') + storage.put_object('bucket', 'stream.bin', io.BytesIO(b'xyz')) + assert storage.get_object('bucket', 'bytes.bin') == b'abc' + assert storage.get_object('bucket', 'text.txt') == b'hello' + stream = storage.get_object('bucket', 'stream.bin', stream=True) + assert stream.read() == b'xyz' + + def test_get_missing_key_raises(self, backend): + storage, _ = backend + with pytest.raises(StorageNoSuchKeyError): + storage.get_object('bucket', 'missing') + + def test_boto3_client_wraps_put_get_and_list(self, backend): + storage, _ = backend + client = storage.get_client() + client.put_object(Bucket='bucket', Key='k', Body=b'v') + assert client.get_object(Bucket='bucket', Key='k')['Body'].read() == b'v' + listed = client.list_objects(Bucket='bucket', Prefix='k') + listed_v2 = client.list_objects_v2(Bucket='bucket', Prefix='k') + assert listed == listed_v2 + assert listed[0]['Key'] == 'k' + assert listed[0]['Size'] == 1 + + def test_upload_and_download_file(self, backend, tmp_path): + storage, _ = backend + src = tmp_path / 'src.txt' + src.write_bytes(b'file-data') + assert storage.upload_file(str(src), 'bucket') is True + dest = tmp_path / 'out' / 'dest.txt' + assert storage.download_file('bucket', 'src.txt', str(dest)) is True + assert dest.read_bytes() == b'file-data' + assert storage.upload_file(str(tmp_path / 'missing.txt'), 'bucket') is False + assert storage.download_file('bucket', 'nope', str(tmp_path / 'x')) is False + + def test_head_object_and_bucket(self, backend): + storage, tmp_path = backend + with pytest.raises(StorageNoSuchKeyError): + storage.head_bucket('bucket') + storage.put_object('bucket', 'dir/key', b'abcd') + assert storage.head_bucket('bucket')['ResponseMetadata']['HTTPStatusCode'] == 200 + assert storage.head_object('bucket', 'dir/key') == {'content-length': '4'} + with pytest.raises(StorageNoSuchKeyError): + storage.head_object('bucket', 'missing') + + def test_delete_object_removes_empty_parents_but_keeps_bucket(self, backend): + storage, tmp_path = backend + storage.put_object('bucket', 'a/b/c.txt', b'x') + storage.delete_object('bucket', 'a/b/c.txt') + bucket_dir = tmp_path / 'bucket' + assert bucket_dir.is_dir() + assert not (bucket_dir / 'a').exists() + + +class TestLocalhostV1Live: + """Live check that version=1 still runs a job end to end.""" + + def test_map_with_localhost_v1(self): + cfg = copy.deepcopy(pytest.lithops_config) + cfg.setdefault('localhost', {}) + cfg['localhost']['version'] = 1 + fexec = lithops.FunctionExecutor(config=cfg) + assert isinstance(fexec.compute_handler, LocalhostHandlerV1) + fexec.map(simple_map_function, [(1, 1)]) + assert fexec.get_result(timeout=20) == [2] + + def test_second_job_survives_the_cleanup_of_the_first(self): + cfg = copy.deepcopy(pytest.lithops_config) + cfg.setdefault('localhost', {}) + cfg['localhost']['version'] = 1 + fexec = lithops.FunctionExecutor(config=cfg) + # v1 runs one job at a time, so the second waits for the first. The + # scoped drain itself is covered by the unit tests: here the window + # where the second job is still queued is too narrow to force + first = fexec.map(sleep_seconds, [1]) + second = fexec.map(simple_map_function, [(2, 2)]) + assert fexec.get_result(fs=first, timeout=20) == [1] + assert fexec.get_result(fs=second, timeout=20) == [4] + + def test_map_with_localhost_v1_multi_worker(self): + cfg = copy.deepcopy(pytest.lithops_config) + cfg.setdefault('localhost', {}) + cfg['localhost']['version'] = 1 + cfg['localhost']['worker_processes'] = 2 + fexec = lithops.FunctionExecutor(config=cfg) + fexec.map(simple_map_function, [(1, 1), (2, 2)]) + assert fexec.get_result(timeout=20) == [2, 4] + + +class TestLocalhostV2Live: + """Live checks that version=2 keeps running jobs across cleanups.""" + + def test_map_after_wait_and_get_result(self): + # Every wait() cleans its jobs up, and a cleanup used to leave + # sentinels in the work queue that killed the next consumers + cfg = copy.deepcopy(pytest.lithops_config) + cfg.setdefault('localhost', {}) + cfg['localhost']['version'] = 2 + fexec = lithops.FunctionExecutor(config=cfg) + assert isinstance(fexec.compute_handler, LocalhostHandlerV2) + fexec.map(simple_map_function, [(1, 1)]) + fexec.wait() + assert fexec.get_result(timeout=20) == [2] + fexec.map(simple_map_function, [(2, 2)]) + assert fexec.get_result(timeout=20) == [4] + + def test_second_job_survives_the_cleanup_of_the_first(self): + cfg = copy.deepcopy(pytest.lithops_config) + cfg.setdefault('localhost', {}) + cfg['localhost']['version'] = 2 + cfg['localhost']['worker_processes'] = 1 + fexec = lithops.FunctionExecutor(config=cfg) + first = fexec.map(simple_map_function, [(1, 1)]) + # More tasks than consumers, so some are still queued when the first + # job is cleaned up: that cleanup must not drop them + second = fexec.map(sleep_seconds, [1, 1, 1]) + assert fexec.get_result(fs=first, timeout=20) == [2] + assert fexec.get_result(fs=second, timeout=30) == [1, 1, 1] + + +def _container_cli(): + return shutil.which('docker') or shutil.which('podman') + + +def _docker_daemon_available(): + cli = _container_cli() + if not cli: + return False + try: + sp.run( + [cli, 'info'], + check=True, + stdout=sp.DEVNULL, + stderr=sp.DEVNULL, + timeout=15, + ) + return True + except Exception: + return False + + +def _python_hub_image(): + return f'python:{sys.version_info.major}.{sys.version_info.minor}' + + +def _ensure_localhost_python_image(): + """Official python:X.Y plus the worker packages Lithops imports at runtime.""" + cli = _container_cli() + base = _python_hub_image() + tag = f'lithops-pytest-{base}' + inspect = sp.run( + [cli, 'image', 'inspect', tag], + stdout=sp.DEVNULL, + stderr=sp.DEVNULL, + ) + if inspect.returncode == 0: + return tag + dockerfile = '\n'.join([ + f'FROM {base}', + 'RUN pip install --no-cache-dir ' + 'cloudpickle tblib pika PyYAML requests tqdm six psutil ps-mem', + ]) + sp.run( + [cli, 'build', '-t', tag, '-'], + input=dockerfile, + check=True, + text=True, + ) + return tag + + +@pytest.mark.skipif( + not _docker_daemon_available(), + reason='docker/podman is not installed or the daemon is not running', +) +class TestLocalhostContainerLive: + """Live localhost jobs inside a Docker Hub python:X.Y container.""" + + def test_map_in_python_container(self): + image = _ensure_localhost_python_image() + cfg = copy.deepcopy(pytest.lithops_config) + cfg.setdefault('lithops', {}) + cfg['lithops']['backend'] = 'localhost' + cfg['lithops']['mode'] = 'localhost' + cfg['lithops']['storage'] = 'localhost' + cfg.setdefault('localhost', {}) + cfg['localhost']['runtime'] = image + cfg['localhost']['pull_runtime'] = False + + with lithops.FunctionExecutor(config=cfg) as fexec: + assert ( + fexec.compute_handler.environment + is localhost_config.LocalhostEnvironment.CONTAINER + ) + fexec.map(simple_map_function, [(2, 3), (4, 5)]) + assert fexec.get_result(timeout=180) == [5, 9] diff --git a/lithops/tests/test_map.py b/lithops/tests/test_map.py index 97d6ab5c2..3ef3a0de4 100644 --- a/lithops/tests/test_map.py +++ b/lithops/tests/test_map.py @@ -25,6 +25,7 @@ lithops_return_futures_call_async, lithops_return_futures_map_multiple, concat, + echo_env_flag, ) @@ -153,3 +154,34 @@ def total(a, b, c=0, *args, d, e=5, **kwargs): result = fexec.get_result() assert result == [22, 22, 20, 30, 32, 32, 32] + + def test_extra_args_tuple_and_dict(self): + fexec = lithops.FunctionExecutor(config=pytest.lithops_config) + fexec.map(simple_map_function, [1, 2, 3], extra_args=(10,)) + assert fexec.get_result() == [11, 12, 13] + + fexec = lithops.FunctionExecutor(config=pytest.lithops_config) + fexec.map( + simple_map_function, + [{'x': 1}, {'x': 2}], + extra_args={'y': 10}, + ) + assert fexec.get_result() == [11, 12] + + def test_extra_env(self): + fexec = lithops.FunctionExecutor(config=pytest.lithops_config) + fexec.map( + echo_env_flag, [1, 2], extra_env={'LITHOPS_TEST_FLAG': 'hello'} + ) + assert fexec.get_result() == ['hello', 'hello'] + + def test_futures_list_chaining(self): + def add_one(x): + return x + 1 + + def mul_two(x): + return x * 2 + + fexec = lithops.FunctionExecutor(config=pytest.lithops_config) + result = fexec.map(add_one, [1, 2, 3]).map(mul_two).get_result() + assert result == [4, 6, 8] diff --git a/lithops/tests/test_monitor.py b/lithops/tests/test_monitor.py new file mode 100644 index 000000000..18d27e0b8 --- /dev/null +++ b/lithops/tests/test_monitor.py @@ -0,0 +1,661 @@ +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import json +import logging +import queue +import threading +import time +from unittest.mock import MagicMock, patch + +import pytest + +from lithops.monitor import ( + LOG_INTERVAL, + JobMonitor, + Monitor, + RabbitmqMonitor, + StorageMonitor, + _is_finished, + _is_started, + _status_id, +) +from lithops import monitor as rabbit_monitor_module +from lithops.utils import ( + _future_id, + monitoring_queue_name, + monitoring_queues, +) + + +class FakeFuture: + def __init__(self, job_id, invoked=False, running=False, ready=False, + success=False, done=False, executor_id='sess-0', call_id='00000', + execution_timeout=10, activation_id=None): + self.job_id = job_id + self.invoked = invoked + self.running = running + self.ready = ready + self.success = success + self.done = done + self.executor_id = executor_id + self.call_id = call_id + self.execution_timeout = execution_timeout + self.activation_id = activation_id + self._call_status = None + self._new_futures = None + self._status_query_count = 0 + + def _set_running(self, call_status): + self._call_status = call_status + self.activation_id = call_status.get('activation_id') + self.running = True + self.invoked = False + + def _set_ready(self, call_status): + self._call_status = call_status + self.ready = True + self.running = False + + def _set_futures(self, call_status): + self._call_status = call_status + self.ready = True + self._new_futures = ['nested'] + + +def _monitor(): + return Monitor( + executor_id='sess-0', + internal_storage=None, + token_bucket_q=None, + job_chunksize={}, + generate_tokens=False, + config={}, + ) + + +class TestMonitorFuturesTracking: + + def test_add_futures_tracks_jobs(self): + monitor = _monitor() + first = FakeFuture('M000') + second = FakeFuture('M000') + third = FakeFuture('M001') + + monitor.add_futures([first, second, third]) + + assert monitor.futures == {first, second, third} + assert monitor.present_jobs == {'M000', 'M001'} + + def test_remove_futures_drops_job_ids_from_removed_set(self): + monitor = _monitor() + keep = FakeFuture('M000', done=True) + drop = FakeFuture('M001', done=True) + monitor.add_futures([keep, drop]) + + monitor.remove_futures([drop]) + + assert monitor.futures == {keep} + assert monitor.present_jobs == {'M000'} + + def test_remove_futures_keeps_job_id_while_siblings_remain(self): + monitor = _monitor() + keep = FakeFuture('M000', done=True) + drop = FakeFuture('M000', done=True) + monitor.add_futures([keep, drop]) + + monitor.remove_futures([drop]) + + assert monitor.futures == {keep} + assert monitor.present_jobs == {'M000'} + + def test_remove_futures_drops_job_id_when_last_sibling_is_removed(self): + monitor = _monitor() + last = FakeFuture('M000', done=True) + other = FakeFuture('M001', done=True) + monitor.add_futures([last, other]) + + monitor.remove_futures([last]) + + assert monitor.futures == {other} + assert monitor.present_jobs == {'M001'} + + def test_all_ready_true_when_every_future_is_terminal(self): + monitor = _monitor() + monitor.add_futures([ + FakeFuture('M000', ready=True), + FakeFuture('M000', success=True), + FakeFuture('M000', done=True), + ]) + assert monitor._all_ready() is True + + def test_all_ready_false_when_any_pending(self): + monitor = _monitor() + monitor.add_futures([ + FakeFuture('M000', ready=True), + FakeFuture('M000'), + ]) + assert monitor._all_ready() is False + + def test_all_ready_swallows_unexpected_errors(self): + class Broken: + @property + def ready(self): + raise RuntimeError('boom') + + monitor = _monitor() + monitor.futures.add(Broken()) + assert monitor._all_ready() is False + + +class TestJobMonitor: + + def test_defaults_to_storage_monitor_without_config(self): + storage = MagicMock() + storage.get_storage_config.return_value = {'monitoring_interval': 2} + storage.backend = 'localhost' + job_monitor = JobMonitor('sess-0', storage) + assert job_monitor.type == 'storage' + assert job_monitor.storage_backend == 'localhost' + + def test_start_creates_monitor_and_records_chunksize(self): + storage = MagicMock() + storage.get_storage_config.return_value = {'monitoring_interval': 2} + storage.backend = 'localhost' + job_monitor = JobMonitor('sess-0', storage) + + instance = MagicMock() + instance.is_alive.return_value = False + job_monitor.MonitorClass = MagicMock(return_value=instance) + + futures = [FakeFuture('M000')] + job_monitor.start(futures, job_id='M000', chunksize=4, generate_tokens=True) + + assert job_monitor.job_chunksize['M000'] == 4 + job_monitor.MonitorClass.assert_called_once() + kwargs = job_monitor.MonitorClass.call_args.kwargs + assert kwargs['generate_tokens'] is True + assert kwargs['config'] == {'monitoring_interval': 2} + instance.add_futures.assert_called_once_with(futures) + instance.start.assert_called_once() + + def test_start_reuses_live_monitor(self): + storage = MagicMock() + storage.get_storage_config.return_value = {'monitoring_interval': 2} + storage.backend = 'localhost' + job_monitor = JobMonitor('sess-0', storage) + + live = MagicMock() + live.is_alive.return_value = True + job_monitor.monitor = live + job_monitor.MonitorClass = MagicMock() + + futures = [FakeFuture('M001')] + job_monitor.start(futures, job_id='M001', chunksize=1) + + job_monitor.MonitorClass.assert_not_called() + live.add_futures.assert_called_once_with(futures) + live.start.assert_not_called() + + def test_rabbitmq_type_uses_backend_section_as_monitor_config(self): + storage = MagicMock() + storage.get_storage_config.return_value = {'monitoring_interval': 2} + storage.backend = 'localhost' + config = { + 'lithops': {'monitoring': 'rabbitmq'}, + 'rabbitmq': {'amqp_url': 'amqp://guest@localhost'}, + } + job_monitor = JobMonitor('sess-0', storage, config=config) + assert job_monitor.type == 'rabbitmq' + assert job_monitor.MonitorClass is RabbitmqMonitor + + instance = MagicMock() + instance.is_alive.return_value = False + job_monitor.MonitorClass = MagicMock(return_value=instance) + job_monitor.start([FakeFuture('M000')]) + assert job_monitor.MonitorClass.call_args.kwargs['config'] == { + 'amqp_url': 'amqp://guest@localhost' + } + + def test_remove_and_stop_are_noops_without_a_live_monitor(self): + storage = MagicMock() + storage.get_storage_config.return_value = {'monitoring_interval': 2} + storage.backend = 'localhost' + job_monitor = JobMonitor('sess-0', storage) + job_monitor.remove([FakeFuture('M000')]) + job_monitor.stop() + + def test_stop_joins_a_live_monitor(self): + storage = MagicMock() + storage.get_storage_config.return_value = {'monitoring_interval': 2} + storage.backend = 'localhost' + job_monitor = JobMonitor('sess-0', storage) + live = MagicMock() + live.is_alive.return_value = True + job_monitor.monitor = live + job_monitor.stop() + live.stop.assert_called_once() + live.join.assert_called_once_with(timeout=5) + + def test_is_alive_requires_a_started_monitor(self): + storage = MagicMock() + storage.get_storage_config.return_value = {'monitoring_interval': 2} + storage.backend = 'localhost' + job_monitor = JobMonitor('sess-0', storage) + with pytest.raises(AttributeError): + job_monitor.is_alive() + + +class TestMonitorHelpers: + + def test_future_and_status_ids_match(self): + future = FakeFuture('M000', call_id='00007') + status = { + 'executor_id': 'sess-0', + 'job_id': 'M000', + 'call_id': '00007', + } + assert _future_id(future) == _status_id(status) == ('sess-0', 'M000', '00007') + + def test_is_finished_and_is_started(self): + pending = FakeFuture('M000', invoked=True) + running = FakeFuture('M000', running=True) + ready = FakeFuture('M000', ready=True) + assert _is_finished(pending) is False + assert _is_started(pending) is False + assert _is_started(running) is True + assert _is_finished(ready) is True + assert _is_started(ready) is True + + +class TestTimeoutAndStatusLog: + + def test_timeout_checker_marks_expired_running_future_ready(self): + monitor = _monitor() + future = FakeFuture('M000', running=True, execution_timeout=1, activation_id='act-1') + future._call_status = {'worker_start_tstamp': time.time() - 100} + monitor._future_timeout_checker([future]) + assert future.ready is True + assert future._call_status['exception'] is True + assert future._call_status['type'] == '__end__' + + def test_timeout_checker_ignores_futures_without_call_status(self): + monitor = _monitor() + future = FakeFuture('M000', running=True) + monitor._future_timeout_checker([future]) + assert future.ready is False + + def test_print_status_log_returns_previous_when_empty(self): + monitor = _monitor() + assert monitor._print_status_log('prev', 3) == ('prev', 3) + + def test_print_status_log_none_log_time_is_short_circuited(self): + monitor = _monitor() + monitor.add_futures([FakeFuture('M000', invoked=True)]) + # Historical: `log_time > LOG_INTERVAL` is not evaluated when counts change. + counts, log_time = monitor._print_status_log(previous_log=None, log_time=None) + assert counts == (1, 0, 0) + assert log_time == 0 + + def test_print_status_log_repeats_after_interval(self): + monitor = _monitor() + monitor.add_futures([FakeFuture('M000', invoked=True)]) + first, _ = monitor._print_status_log(previous_log=None, log_time=0) + same, log_time = monitor._print_status_log(previous_log=first, log_time=0) + assert log_time == 0 + _, log_time = monitor._print_status_log( + previous_log=first, log_time=LOG_INTERVAL + 1 + ) + assert log_time == 0 + + def test_print_status_log_does_not_repeat_when_all_finished(self, caplog): + monitor = _monitor() + monitor.add_futures([FakeFuture('M000', invoked=True, ready=True)]) + first, _ = monitor._print_status_log(previous_log=None, log_time=0) + with caplog.at_level(logging.DEBUG, logger='lithops.monitor'): + counts, log_time = monitor._print_status_log( + previous_log=first, log_time=LOG_INTERVAL + 1 + ) + assert counts == first + assert log_time == LOG_INTERVAL + 1 + assert caplog.records == [] + + def test_check_new_futures_updates_tracking_set(self): + monitor = _monitor() + future = FakeFuture('M000') + monitor.add_futures([future]) + assert monitor._check_new_futures({'type': '__end__'}, future) is False + assert monitor._check_new_futures({'new_futures': 'x'}, future) is True + assert 'nested' in monitor.futures + + +class TestStorageMonitorTokensAndTags: + + def _storage(self, generate_tokens=True, chunksize=2): + storage = MagicMock() + q = queue.Queue() + return StorageMonitor( + executor_id='sess-0', + internal_storage=storage, + token_bucket_q=q, + job_chunksize={'M000': chunksize}, + generate_tokens=generate_tokens, + config={'monitoring_interval': 1}, + ) + + def test_generate_tokens_skips_when_disabled(self): + monitor = self._storage(generate_tokens=False) + monitor._generate_tokens({(('sess-0', 'M000', '00000'), 'w1')}, set()) + assert monitor.token_bucket_q.empty() + + def test_generate_tokens_emits_when_chunk_completes(self): + monitor = self._storage() + running = { + (('sess-0', 'M000', '00000'), 'w1'), + (('sess-0', 'M000', '00001'), 'w1'), + } + done = {('sess-0', 'M000', '00000'), ('sess-0', 'M000', '00001')} + monitor.present_jobs.add('M000') + monitor._generate_tokens(running, done) + assert monitor.token_bucket_q.get_nowait() == '#' + assert 'w1' in monitor.workers_done + + def test_generate_tokens_waits_for_full_chunk_then_does_not_repeat(self): + monitor = self._storage() + running = { + (('sess-0', 'M000', '00000'), 'w1'), + (('sess-0', 'M000', '00001'), 'w1'), + } + first_done = {('sess-0', 'M000', '00000')} + all_done = first_done | {('sess-0', 'M000', '00001')} + monitor.present_jobs.add('M000') + monitor._generate_tokens(running, first_done) + assert monitor.token_bucket_q.empty() + monitor._generate_tokens(running, all_done) + assert monitor.token_bucket_q.get_nowait() == '#' + monitor._generate_tokens(running, all_done) + assert monitor.token_bucket_q.empty() + + def test_generate_tokens_one_per_worker(self): + monitor = self._storage(chunksize=1) + running = { + (('sess-0', 'M000', '00000'), 'w1'), + (('sess-0', 'M000', '00001'), 'w2'), + } + done = {('sess-0', 'M000', '00000'), ('sess-0', 'M000', '00001')} + monitor.present_jobs.add('M000') + monitor._generate_tokens(running, done) + tokens = [ + monitor.token_bucket_q.get_nowait(), + monitor.token_bucket_q.get_nowait(), + ] + assert tokens == ['#', '#'] + assert monitor.token_bucket_q.empty() + + def test_generate_tokens_skips_job_not_present(self): + monitor = self._storage() + running = { + (('sess-0', 'M000', '00000'), 'w1'), + (('sess-0', 'M000', '00001'), 'w1'), + } + done = {('sess-0', 'M000', '00000'), ('sess-0', 'M000', '00001')} + monitor._generate_tokens(running, done) + assert monitor.token_bucket_q.empty() + + def test_tag_future_as_running_from_callids(self): + monitor = self._storage() + future = FakeFuture('M000', invoked=True, call_id='00000') + monitor.add_futures([future]) + callids_running = {(('sess-0', 'M000', '00000'), 'act-9')} + monitor._tag_future_as_running(callids_running) + assert future.running is True + assert future.activation_id == 'act-9' + + def test_tag_future_as_ready_queries_matching_ids(self): + monitor = self._storage() + future = FakeFuture('M000', invoked=True, call_id='00000') + monitor.add_futures([future]) + monitor.internal_storage.get_call_status.return_value = { + 'type': '__end__', + 'activation_id': 'act-9', + } + monitor._tag_future_as_ready({('sess-0', 'M000', '00000')}) + assert future.ready is True + assert future._call_status['activation_id'] == 'act-9' + + def test_tag_future_as_ready_queries_only_matching_ids_when_not_near_complete( + self, + ): + monitor = self._storage() + futures = [ + FakeFuture('M000', invoked=True, call_id=f'{i:05d}') + for i in range(20) + ] + monitor.add_futures(futures) + monitor.internal_storage.get_call_status.return_value = { + 'type': '__end__', + 'activation_id': 'act', + } + monitor._tag_future_as_ready({('sess-0', 'M000', '00003')}) + queried = [ + call.args for call in monitor.internal_storage.get_call_status.call_args_list + ] + assert queried == [('sess-0', 'M000', '00003')] + assert futures[3].ready is True + assert futures[0].ready is False + + def test_poll_and_process_returns_new_done_ids_and_tags(self): + monitor = self._storage() + monitor._generate_tokens = MagicMock() + monitor._tag_future_as_running = MagicMock() + monitor._tag_future_as_ready = MagicMock() + monitor._print_status_log = MagicMock(return_value=('log', 1)) + running = {(('sess-0', 'M000', '00000'), 'w1')} + done = {('sess-0', 'M000', '00000')} + monitor.internal_storage.get_job_status.return_value = (running, done) + new, prev, log_time = monitor._poll_and_process_job_status(None, 0) + assert new == done + monitor.internal_storage.get_job_status.assert_called_once_with( + 'sess-0', job_ids=set() + ) + monitor._generate_tokens.assert_called_once_with(running, done) + monitor._tag_future_as_running.assert_called_once_with(running) + monitor._tag_future_as_ready.assert_called_once_with(done) + assert prev == 'log' + assert log_time == 1 + + def test_poll_and_process_emits_token_when_chunk_completes(self): + monitor = self._storage() + future0 = FakeFuture('M000', invoked=True, call_id='00000') + future1 = FakeFuture('M000', invoked=True, call_id='00001') + monitor.add_futures([future0, future1]) + running = { + (('sess-0', 'M000', '00000'), 'w1'), + (('sess-0', 'M000', '00001'), 'w1'), + } + done = {('sess-0', 'M000', '00000'), ('sess-0', 'M000', '00001')} + monitor.internal_storage.get_job_status.return_value = (running, done) + monitor.internal_storage.get_call_status.return_value = { + 'type': '__end__', + 'activation_id': 'w1', + } + monitor._print_status_log = MagicMock(return_value=('log', 1)) + monitor._poll_and_process_job_status(None, 0) + assert monitor.token_bucket_q.get_nowait() == '#' + assert future0.ready is True + assert future1.ready is True + + def test_run_sleeps_shorter_when_new_done_then_polls_after_loop(self): + monitor = self._storage() + polls = [] + + def poll(previous_log, log_time): + polls.append(1) + if len(polls) == 1: + return {('sess-0', 'M000', '00000')}, previous_log, log_time + monitor.should_run = False + return set(), previous_log, log_time + + monitor._poll_and_process_job_status = poll + sleeps = [] + test_thread = threading.current_thread() + + def sleep(seconds): + if threading.current_thread() is test_thread: + sleeps.append(seconds) + + with patch('lithops.monitor.time.sleep', side_effect=sleep): + monitor.run() + assert sleeps == [0.2] + assert len(polls) == 3 + + def test_run_skips_sleep_and_swallows_final_poll_errors_after_stop(self): + monitor = self._storage() + polls = [] + + def poll(previous_log, log_time): + polls.append(1) + monitor.should_run = False + if len(polls) > 1: + raise RuntimeError('storage gone') + return set(), previous_log, log_time + + monitor._poll_and_process_job_status = poll + sleeps = [] + test_thread = threading.current_thread() + + def sleep(seconds): + if threading.current_thread() is test_thread: + sleeps.append(seconds) + + with patch('lithops.monitor.time.sleep', side_effect=sleep): + monitor.run() + assert sleeps == [] + assert len(polls) == 2 + + +class TestRabbitmqMonitorTags: + + def _rabbit(self): + monitor = RabbitmqMonitor.__new__(RabbitmqMonitor) + Monitor.__init__( + monitor, 'sess-0', None, queue.Queue(), {'M000': 1}, True, {} + ) + return monitor + + def test_the_declared_queue_is_the_one_the_workers_publish_to(self): + # The monitor declares and consumes one queue; the workers derive the + # names they publish to from the same helper. If these two ever drift + # apart, every rabbitmq job hangs without a word + pika = rabbit_monitor_module.pika + with patch.object(pika, 'URLParameters'), \ + patch.object(pika, 'BlockingConnection') as connection: + monitor = RabbitmqMonitor( + 'sess-0', None, queue.Queue(), {'M000': 1}, True, + {'amqp_url': 'amqp://guest:guest@localhost:5672'}, + ) + + assert monitor.queue == monitoring_queue_name('sess-0') + assert monitoring_queues('sess-0') == [monitor.queue] + declared = connection.return_value.channel.return_value.queue_declare + assert declared.call_args.kwargs['queue'] == monitor.queue + + def test_tag_running_and_ready_by_call_status(self): + monitor = self._rabbit() + future = FakeFuture('M000', invoked=True, call_id='00000') + monitor.add_futures([future]) + init = { + 'type': '__init__', + 'executor_id': 'sess-0', + 'job_id': 'M000', + 'call_id': '00000', + 'activation_id': 'act-1', + } + monitor._tag_future_as_running(init) + assert future.running is True + end = { + 'type': '__end__', + 'executor_id': 'sess-0', + 'job_id': 'M000', + 'call_id': '00000', + 'activation_id': 'act-1', + 'chunksize': 1, + } + monitor._tag_future_as_ready(end) + assert future.ready is True + + def test_generate_tokens_emits_after_chunksize_completions(self): + monitor = self._rabbit() + status = { + 'activation_id': 'w1', + 'executor_id': 'sess-0', + 'job_id': 'M000', + 'call_id': '00000', + 'chunksize': 1, + } + monitor._generate_tokens(status) + assert monitor.token_bucket_q.get_nowait() == '#' + + def test_generate_tokens_waits_for_chunksize_completions(self): + monitor = self._rabbit() + first = { + 'activation_id': 'w1', + 'executor_id': 'sess-0', + 'job_id': 'M000', + 'call_id': '00000', + 'chunksize': 2, + } + monitor._generate_tokens(first) + assert monitor.token_bucket_q.empty() + second = dict(first, call_id='00001') + monitor._generate_tokens(second) + assert monitor.token_bucket_q.get_nowait() == '#' + + def test_run_processes_init_and_end_until_all_ready(self): + monitor = self._rabbit() + future = FakeFuture('M000', invoked=True, call_id='00000') + monitor.add_futures([future]) + monitor.queue = 'lithops-sess-0' + monitor.should_run = True + monitor._print_status_log = MagicMock(return_value=(None, 0)) + + channel = MagicMock() + monitor.connection = MagicMock() + monitor.connection.channel.return_value = channel + + def consume(): + callback = channel.basic_consume.call_args[0][1] + init = json.dumps({ + 'type': '__init__', + 'executor_id': 'sess-0', + 'job_id': 'M000', + 'call_id': '00000', + 'activation_id': 'act-1', + }).encode() + end = json.dumps({ + 'type': '__end__', + 'executor_id': 'sess-0', + 'job_id': 'M000', + 'call_id': '00000', + 'activation_id': 'act-1', + 'chunksize': 1, + }).encode() + callback(channel, None, None, init) + callback(channel, None, None, end) + + channel.start_consuming.side_effect = consume + with patch('lithops.monitor.threading.Thread'): + monitor.run() + assert future.ready is True + channel.stop_consuming.assert_called() diff --git a/lithops/tests/test_plots.py b/lithops/tests/test_plots.py new file mode 100644 index 000000000..297bf393a --- /dev/null +++ b/lithops/tests/test_plots.py @@ -0,0 +1,136 @@ +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import os +from unittest.mock import MagicMock + +import pytest + +seaborn = pytest.importorskip('seaborn') # noqa: F401 + +from lithops.plots import ( # noqa: E402 + _elapsed, + _plot_destination, + _set_call_axis, + _set_time_axis, + _timeline_span, + create_histogram, + create_timeline, +) + + +class TestPlotDestination: + + def test_none_writes_under_plots_with_suffix(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + monkeypatch.setattr('lithops.plots.time.time', lambda: 1234) + path = _plot_destination(None, 'timeline.png') + assert path == os.path.join(str(tmp_path), 'plots', '1234_timeline.png') + assert os.path.isdir(tmp_path / 'plots') + + def test_explicit_path_appends_suffix(self, tmp_path): + dst = tmp_path / 'run' + dst.write_text('') + path = _plot_destination(str(dst), 'histogram.png') + assert path == '{}_{}'.format(os.path.realpath(dst), 'histogram.png') + + def test_expands_user_home_when_tilde_present(self, tmp_path, monkeypatch): + monkeypatch.setenv('HOME', str(tmp_path)) + path = _plot_destination(os.path.join('~', 'out'), 'timeline.png') + expected = '{}_{}'.format( + os.path.realpath(os.path.join(str(tmp_path), 'out')), + 'timeline.png', + ) + assert path == expected + + +class TestPlotAxes: + + def test_call_axis_sets_ticks_and_ylim(self): + ax = MagicMock() + y_ticks = _set_call_axis(ax, 40) + ax.set_yticks.assert_called_once() + ax.set_ylim.assert_called_once() + assert len(y_ticks) > 0 + + def test_time_axis_draws_vertical_guides(self): + ax = MagicMock() + x_ticks = _set_time_axis(ax, 16) + ax.set_xlim.assert_called_once_with(0, 16) + ax.set_xticks.assert_called_once() + assert ax.axvline.call_count == len(x_ticks) + + +class TestElapsedAndSpan: + + def test_elapsed_subtracts_origin(self): + import pandas as pd + series = pd.Series([10.0, 12.0]) + assert list(_elapsed(series, 10.0)) == [0.0, 2.0] + + def test_timeline_span_prefers_result_then_status_then_end(self): + import pandas as pd + t0 = 100.0 + with_result = pd.DataFrame({ + 'host_result_done_tstamp': [104.0], + 'host_status_done_tstamp': [103.0], + 'end_tstamp': [102.0], + }) + assert _timeline_span(with_result, t0) == pytest.approx(4.0 * 1.25) + with_status = pd.DataFrame({ + 'host_status_done_tstamp': [108.0], + 'end_tstamp': [102.0], + }) + assert _timeline_span(with_status, t0) == pytest.approx(8.0 * 1.25) + with_end = pd.DataFrame({'end_tstamp': [110.0]}) + assert _timeline_span(with_end, t0) == pytest.approx(10.0 * 1.25) + + +def _stats(t0=1000.0, with_result=True): + stats = { + 'host_job_create_tstamp': t0, + 'host_submit_tstamp': t0 + 0.1, + 'worker_func_start_tstamp': t0 + 0.2, + 'worker_func_end_tstamp': t0 + 0.8, + 'host_status_done_tstamp': t0 + 0.9, + 'worker_start_tstamp': t0 + 0.2, + 'worker_end_tstamp': t0 + 0.8, + } + if with_result: + stats['host_result_done_tstamp'] = t0 + 1.0 + return stats + + +class FakePlotFuture: + def __init__(self, stats): + self.stats = stats + + +class TestCreatePlots: + + def test_create_timeline_and_histogram_write_files(self, tmp_path): + fs = [FakePlotFuture(_stats()), FakePlotFuture(_stats(t0=1000.2))] + dest = str(tmp_path / 'run') + create_timeline(fs, dest, figsize=(4, 3)) + create_histogram(fs, dest, figsize=(4, 3)) + timeline = '{}_{}'.format(os.path.realpath(dest), 'timeline.png') + histogram = '{}_{}'.format(os.path.realpath(dest), 'histogram.png') + assert os.path.isfile(timeline) + assert os.path.isfile(histogram) + + def test_create_timeline_without_result_timestamps(self, tmp_path): + fs = [FakePlotFuture(_stats(with_result=False))] + dest = str(tmp_path / 'status-only') + create_timeline(fs, dest, figsize=(4, 3)) + assert os.path.isfile('{}_{}'.format(os.path.realpath(dest), 'timeline.png')) diff --git a/lithops/tests/test_retries.py b/lithops/tests/test_retries.py index e952c8c13..813bc1186 100644 --- a/lithops/tests/test_retries.py +++ b/lithops/tests/test_retries.py @@ -1,8 +1,12 @@ import time +from unittest.mock import MagicMock + import pytest from lithops import FunctionExecutor from lithops import RetryingFunctionExecutor +from lithops.retries import RetryingFuture +from lithops.wait import ALWAYS, ANY_COMPLETED def run_test(function, input, retries, timeout=5): @@ -154,3 +158,192 @@ def check_invocation_counts( f"Invocation count for {i}, expected: {expected_count}, actual: {actual_count}" ) assert actual_invocation_counts == expected_invocation_counts + + +class FakeResponseFuture: + def __init__(self, error=False, result=None, done=False): + self.error = error + self.done = done + self._result = result + self._status = 'ok' + self._exception = (RuntimeError, RuntimeError('failed'), None) + self.stats = {'worker_exec_time': 1.0} + + def status(self, throw_except=True, internal_storage=None, check_only=False): + return self._status + + def result(self, throw_except=True, internal_storage=None): + return self._result + + +class TestRetryingFutureUnit: + + def test_should_retry_until_budget_exhausted(self): + future = RetryingFuture(FakeResponseFuture(), map_function=lambda x: x, input=1, retries=1) + assert future._should_retry() is True + future._inc_failure_count() + assert future.failure_count == 1 + assert future._should_retry() is True + future._inc_failure_count() + assert future._should_retry() is False + + def test_cancel_prevents_retry(self): + future = RetryingFuture(FakeResponseFuture(), map_function=lambda x: x, input=1, retries=5) + future.cancel() + assert future._should_retry() is False + + def test_retries_default_to_zero(self): + future = RetryingFuture(FakeResponseFuture(), map_function=lambda x: x, input=1) + assert future.retries == 0 + + def test_status_and_result_reraise_on_error(self): + wrapped = FakeResponseFuture(error=True, result='nope') + future = RetryingFuture(wrapped, map_function=lambda x: x, input=1, retries=0) + with pytest.raises(RuntimeError, match='failed'): + future.status() + with pytest.raises(RuntimeError, match='failed'): + future.result() + + def test_status_and_result_passthrough_on_success(self): + wrapped = FakeResponseFuture(error=False, result=42) + future = RetryingFuture(wrapped, map_function=lambda x: x, input=1) + assert future.status() == 'ok' + assert future.result() == 42 + assert future.done is False + assert future.stats == wrapped.stats + + def test_retry_resubmits_original_input_and_kwargs(self): + replacement = FakeResponseFuture(result=99) + executor = MagicMock() + executor.map.return_value = [replacement] + future = RetryingFuture( + FakeResponseFuture(error=True), + map_function=str, + input=7, + retries=1, + timeout=11, + ) + future._retry(executor) + executor.map.assert_called_once_with(str, [7], timeout=11) + assert future.response_future is replacement + + +class TestRetryingFunctionExecutorUnit: + + def test_map_uses_config_retries_and_forwards_kwargs(self): + inner = MagicMock() + inner.config = {'lithops': {'retries': 4}} + inner.map.return_value = [FakeResponseFuture(), FakeResponseFuture()] + executor = RetryingFunctionExecutor(inner) + + futures = executor.map( + lambda x: x, + [1, 2], + timeout=9, + extra_env={'A': '1'}, + chunksize=3, + obj_chunk_size=10, + obj_chunk_number=2, + obj_newline=None, + ) + + assert [f.retries for f in futures] == [4, 4] + assert [f.input for f in futures] == [1, 2] + inner.map.assert_called_once() + kwargs = inner.map.call_args.kwargs + assert kwargs['timeout'] == 9 + assert kwargs['extra_env'] == {'A': '1'} + assert kwargs['chunksize'] == 3 + assert kwargs['obj_chunk_size'] == 10 + assert kwargs['obj_chunk_number'] == 2 + assert kwargs['obj_newline'] is None + assert futures[0].map_kwargs['timeout'] == 9 + + def test_explicit_retries_override_config(self): + inner = MagicMock() + inner.config = {'lithops': {'retries': 4}} + inner.map.return_value = [FakeResponseFuture()] + executor = RetryingFunctionExecutor(inner) + + futures = executor.map(lambda x: x, [1], retries=0) + assert futures[0].retries == 0 + + def test_wait_retries_failed_futures_until_all_complete(self): + first = FakeResponseFuture(error=True) + retried = FakeResponseFuture(error=False, result=1) + inner = MagicMock() + inner.config = {} + inner.wait.side_effect = [ + ([first], []), + ([retried], []), + ] + inner.map.return_value = [retried] + + retrying = RetryingFuture(first, map_function=lambda x: x, input=1, retries=1) + executor = RetryingFunctionExecutor(inner) + done, pending = executor.wait([retrying], throw_except=False) + + assert pending == [] + assert done == [retrying] + assert retrying.response_future is retried + inner.map.assert_called_once() + + def test_wait_always_returns_after_first_poll(self): + pending_resp = FakeResponseFuture() + inner = MagicMock() + inner.config = {} + inner.wait.return_value = ([], [pending_resp]) + retrying = RetryingFuture(pending_resp, map_function=lambda x: x, input=1, retries=0) + executor = RetryingFunctionExecutor(inner) + + done, pending = executor.wait([retrying], return_when=ALWAYS) + assert done == [] + assert pending == [retrying] + assert inner.wait.call_count == 1 + + def test_wait_any_completed_stops_when_one_succeeds(self): + finished = FakeResponseFuture(error=False, result=1) + pending = FakeResponseFuture() + inner = MagicMock() + inner.config = {} + inner.wait.return_value = ([finished], [pending]) + done_f = RetryingFuture(finished, map_function=lambda x: x, input=1, retries=0) + pending_f = RetryingFuture(pending, map_function=lambda x: x, input=2, retries=0) + executor = RetryingFunctionExecutor(inner) + done, still_pending = executor.wait( + [done_f, pending_f], return_when=ANY_COMPLETED + ) + assert done == [done_f] + assert still_pending == [pending_f] + assert inner.wait.call_count == 1 + + def test_wait_exhausted_retries_are_treated_as_done(self): + failed = FakeResponseFuture(error=True) + inner = MagicMock() + inner.config = {} + inner.wait.return_value = ([failed], []) + retrying = RetryingFuture(failed, map_function=lambda x: x, input=1, retries=0) + executor = RetryingFunctionExecutor(inner) + done, pending = executor.wait([retrying], throw_except=False) + assert pending == [] + assert done == [retrying] + inner.map.assert_not_called() + + def test_context_manager_and_clean_delegate(self): + inner = MagicMock() + executor = RetryingFunctionExecutor(inner) + with executor: + pass + inner.__enter__.assert_called_once() + inner.__exit__.assert_called_once() + executor.clean('fs', 'cs', False, True, True) + inner.clean.assert_called_once_with('fs', 'cs', False, True, True) + + def test_retries_to_use_prefers_explicit_then_config(self): + inner = MagicMock() + inner.config = {'lithops': {'retries': 9}} + executor = RetryingFunctionExecutor(inner) + assert executor._retries_to_use(3) == 3 + assert executor._retries_to_use(None) == 9 + executor.config = {} + assert executor._retries_to_use(None) == 0 diff --git a/lithops/tests/test_scripts.py b/lithops/tests/test_scripts.py new file mode 100644 index 000000000..0128e5bd2 --- /dev/null +++ b/lithops/tests/test_scripts.py @@ -0,0 +1,574 @@ +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import os +import pickle +import sys +from datetime import datetime +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +from click.testing import CliRunner + +from lithops.constants import ( + CACHE_DIR, + JOBS_PREFIX, + LOCALHOST, + LITHOPS_TEMP_DIR, + RUNTIMES_PREFIX, + SERVERLESS, + STANDALONE, + TEMP_PREFIX, +) +from lithops.scripts import cleaner +from lithops.scripts import cli as cli_module +from lithops.scripts.cli import ( + _compute_handler, + _format_storage_objects, + _localize_and_sort_rows, + _require_mode, + lithops_cli, + set_config_ow, +) +from lithops.scripts.cleaner import ( + _classify_cleaner_files, + _executor_id_from_jobs, + _run_clean_tasks, + clean_cloudobjects, + clean_executor_jobs, + clean_functions, +) + + +def _cli(*args): + return CliRunner().invoke(lithops_cli, list(args)) + + +class TestSetConfigOw: + + def test_empty(self): + assert set_config_ow() == {'lithops': {}, 'backend': {}} + + def test_backend_sets_mode(self): + cfg = set_config_ow(backend=LOCALHOST) + assert cfg['lithops']['backend'] == LOCALHOST + assert cfg['lithops']['mode'] == LOCALHOST + + def test_optional_fields(self): + cfg = set_config_ow( + storage='ibm_cos', runtime_name='rt', region='eu' + ) + assert cfg['lithops']['storage'] == 'ibm_cos' + assert cfg['backend']['runtime'] == 'rt' + assert cfg['backend']['region'] == 'eu' + + +class TestCliHelpers: + + def test_require_mode_standalone_mentions_command(self): + with pytest.raises(Exception, match='lithops image list') as exc: + _require_mode( + {'lithops': {'mode': LOCALHOST}}, + STANDALONE, + 'lithops image list', + ) + assert 'image build' not in str(exc.value) + + def test_require_mode_serverless(self): + with pytest.raises(Exception, match='serverless'): + _require_mode( + {'lithops': {'mode': LOCALHOST}}, + SERVERLESS, + 'lithops runtime build', + ) + + def test_require_mode_ok(self): + _require_mode( + {'lithops': {'mode': STANDALONE}}, STANDALONE, 'lithops job list' + ) + + def test_compute_handler_unknown_mode_raises(self): + with pytest.raises(Exception, match='Unknown compute mode'): + _compute_handler({'lithops': {'mode': 'nope'}}) + + def test_format_storage_objects(self): + modified = datetime(2024, 1, 2, 3, 4, 5) + rows = _format_storage_objects([ + {'Key': 'a', 'LastModified': modified, 'Size': 1024}, + {'Key': 'b'}, + ]) + assert rows[0]['Key'] == 'a' + assert 'Jan' in rows[0]['LastModified'] + assert rows[0]['Size'] + assert rows[1] == {'Key': 'b'} + + def test_format_storage_objects_empty(self): + assert _format_storage_objects([]) == [] + + def test_localize_and_sort_rows(self): + rows = [ + ['b', '2024-01-02 00:00:00 UTC'], + ['a', '2024-01-01 00:00:00 UTC'], + ] + sorted_rows = _localize_and_sort_rows(rows, 1) + assert sorted_rows[0][0] == 'a' + + +class TestStorageCommands: + + def test_list_empty_bucket_does_not_indexerror(self): + client = MagicMock() + client.backend = 'localhost' + client.list_objects.return_value = [] + with patch('lithops.scripts.cli.Storage', return_value=client): + with patch('lithops.scripts.cli.setup_lithops_logger'): + result = _cli('storage', 'list', 'bucket') + assert result.exit_code == 0 + assert result.exception is None + assert 'No information' in result.output + + def test_list_objects_prints_table(self): + client = MagicMock() + client.backend = 'localhost' + client.list_objects.return_value = [ + {'Key': 'a.txt', 'Size': 10, 'LastModified': datetime(2024, 1, 1)}, + ] + with patch('lithops.scripts.cli.Storage', return_value=client): + with patch('lithops.scripts.cli.setup_lithops_logger'): + result = _cli('storage', 'list', 'bucket') + assert result.exit_code == 0 + assert 'a.txt' in result.output + assert 'Total objects: 1' in result.output + + def test_delete_requires_key_or_prefix(self): + with patch('lithops.scripts.cli.Storage', return_value=MagicMock()): + with patch('lithops.scripts.cli.setup_lithops_logger'): + result = _cli('storage', 'delete', 'bucket') + assert result.exit_code != 0 + assert 'KEY or --prefix' in result.output + + def test_delete_key(self): + client = MagicMock() + with patch('lithops.scripts.cli.Storage', return_value=client): + with patch('lithops.scripts.cli.setup_lithops_logger'): + result = _cli('storage', 'delete', 'bucket', 'obj') + assert result.exit_code == 0 + client.delete_object.assert_called_once_with('bucket', 'obj') + + def test_delete_prefix(self): + client = MagicMock() + client.list_keys.return_value = ['a', 'b'] + with patch('lithops.scripts.cli.Storage', return_value=client): + with patch('lithops.scripts.cli.setup_lithops_logger'): + result = _cli('storage', 'delete', 'bucket', '--prefix', 'pre') + assert result.exit_code == 0 + client.delete_objects.assert_called_once_with('bucket', ['a', 'b']) + + +class TestLogsAndAttach: + + def test_get_logs_missing_file_message(self, tmp_path, monkeypatch): + monkeypatch.setattr('lithops.scripts.cli.LOGS_DIR', str(tmp_path)) + result = _cli('logs', 'get', 'missing-id') + assert result.exit_code == 0 + assert 'does not exist' in result.output + assert 'does not exists' not in result.output + + def test_get_logs_prints_file(self, tmp_path, monkeypatch): + monkeypatch.setattr('lithops.scripts.cli.LOGS_DIR', str(tmp_path)) + (tmp_path / 'job.log').write_text('hello-log\n') + result = _cli('logs', 'get', 'job') + assert 'hello-log' in result.output + + def test_attach_ssh_uses_argv_list_for_key_with_spaces(self): + handler = MagicMock() + handler.is_initialized.return_value = True + handler.backend.master.is_ready.return_value = True + handler.backend.master.get_public_ip.return_value = '10.0.0.1' + key = '/tmp/my key.pem' + handler.backend.master.ssh_credentials = { + 'username': 'ubuntu', + 'key_filename': key, + } + with patch('lithops.scripts.cli._prepare_standalone', return_value=handler): + with patch('lithops.scripts.cli.os.path.exists', return_value=True): + with patch('lithops.scripts.cli.sp.run') as run: + result = _cli('attach', '-b', 'aws_ec2') + assert result.exit_code == 0 + cmd = run.call_args[0][0] + assert cmd[0] == 'ssh' + assert isinstance(cmd, list) + assert os.path.abspath(os.path.expanduser(key)) in cmd + assert 'ubuntu@10.0.0.1' in cmd + + +class TestHelloAndClean: + + def test_hello_call_async(self): + fexec = MagicMock() + fexec.get_result.return_value = 'Hello tester!' + with patch('getpass.getuser', return_value='tester'): + with patch( + 'lithops.scripts.cli.lithops.FunctionExecutor', + return_value=fexec, + ): + with patch('lithops.scripts.cli.setup_lithops_logger'): + result = _cli('hello') + assert result.exit_code == 0 + fexec.call_async.assert_called_once() + assert 'Lithops is working as expected' in result.output + + def test_clean_localhost(self): + cfg = { + 'lithops': {'mode': LOCALHOST, 'backend': LOCALHOST}, + LOCALHOST: {}, + } + handler = MagicMock() + storage = MagicMock() + storage.bucket = 'bkt' + internal = MagicMock() + internal.storage = storage + with patch('lithops.scripts.cli._resolved_config', return_value=cfg): + with patch('lithops.scripts.cli.extract_storage_config', return_value={}): + with patch('lithops.scripts.cli.InternalStorage', return_value=internal): + with patch( + 'lithops.scripts.cli._compute_handler', + return_value=handler, + ): + with patch('lithops.scripts.cli.clean_bucket') as cb: + with patch('lithops.scripts.cli.shutil.rmtree') as rmtree: + with patch( + 'lithops.scripts.cli._clean_local_temp_data' + ) as clean_temp: + with patch('lithops.scripts.cli.setup_lithops_logger'): + result = _cli('clean', '--all') + assert result.exit_code == 0 + handler.clean.assert_called_once_with(all=True) + assert cb.call_count == 2 + clean_temp.assert_called_once_with() + removed = [call.args[0] for call in rmtree.call_args_list] + assert LITHOPS_TEMP_DIR not in removed + assert os.path.join(CACHE_DIR, RUNTIMES_PREFIX, LOCALHOST) in removed + + def test_clean_local_temp_keeps_the_shared_skeleton( + self, tmp_path, monkeypatch + ): + cleaner_dir = tmp_path / 'cleaner' + logs_dir = tmp_path / 'logs' + jobs_dir = tmp_path / 'jobs' + for path in (cleaner_dir, logs_dir, jobs_dir, tmp_path / 'modules'): + path.mkdir() + (cleaner_dir / 'pending-request').write_text('keep me') + (logs_dir / 'ex-0-A000.log').write_text('drop me') + (tmp_path / 'functions.log').write_text('drop me') + + for name, value in ( + ('LITHOPS_TEMP_DIR', tmp_path), + ('CLEANER_DIR', cleaner_dir), + ('LOGS_DIR', logs_dir), + ('JOBS_DIR', jobs_dir), + ): + monkeypatch.setattr(f'lithops.scripts.cli.{name}', str(value)) + + cli_module._clean_local_temp_data() + + # The requests of the other processes on this machine survive + assert (cleaner_dir / 'pending-request').read_text() == 'keep me' + # Their data does not, but the directories they write into come back + assert not (logs_dir / 'ex-0-A000.log').exists() + assert not (tmp_path / 'functions.log').exists() + assert not (tmp_path / 'modules').exists() + assert logs_dir.is_dir() + assert jobs_dir.is_dir() + + def test_clean_local_temp_tolerates_a_missing_dir(self, tmp_path, monkeypatch): + missing = tmp_path / 'gone' + monkeypatch.setattr('lithops.scripts.cli.LITHOPS_TEMP_DIR', str(missing)) + for name in ('CLEANER_DIR', 'LOGS_DIR', 'JOBS_DIR'): + monkeypatch.setattr( + f'lithops.scripts.cli.{name}', str(missing / name.lower()) + ) + + cli_module._clean_local_temp_data() + + assert missing.is_dir() + + +class TestJobWorkerList: + + def test_job_list_empty_does_not_crash(self): + handler = MagicMock() + handler.is_initialized.return_value = True + handler.backend.master.is_ready.return_value = True + handler._is_master_service_ready.return_value = True + handler.list_jobs.return_value = [] + with patch('lithops.scripts.cli._prepare_standalone', return_value=handler): + result = _cli('job', 'list', '-b', 'aws_ec2') + assert result.exit_code == 0 + assert 'Total jobs: 0' in result.output + + def test_worker_list_empty_does_not_crash(self): + handler = MagicMock() + handler.is_initialized.return_value = True + handler.backend.master.is_ready.return_value = True + handler._is_master_service_ready.return_value = True + handler.list_workers.return_value = [] + with patch('lithops.scripts.cli._prepare_standalone', return_value=handler): + result = _cli('worker', 'list', '-b', 'aws_ec2') + assert result.exit_code == 0 + assert 'Total workers: 0' in result.output + + def test_standalone_not_initialized(self): + handler = MagicMock() + handler.is_initialized.return_value = False + with patch('lithops.scripts.cli._prepare_standalone', return_value=handler): + result = _cli('job', 'list') + assert result.exit_code == 0 + handler.list_jobs.assert_not_called() + + +class TestCleaner: + + def test_import_does_not_redirect_stdout(self): + dest = getattr(sys.stdout, 'name', None) + assert dest != cleaner.CLEANER_LOG_FILE + + def test_executor_id_from_jobs(self): + assert _executor_id_from_jobs({'abc-0-M000'}) == 'abc-0' + assert _executor_id_from_jobs(set()) is None + + def test_empty_jobs_to_clean_is_skipped(self, tmp_path, monkeypatch): + monkeypatch.setattr(cleaner, 'CLEANER_DIR', str(tmp_path)) + payload = { + 'jobs_to_clean': set(), + 'storage_config': {}, + 'clean_cloudobjects': False, + } + path = tmp_path / 'job.pkl' + path.write_bytes(pickle.dumps(payload)) + jobs, cos, fns = _classify_cleaner_files(['job.pkl']) + assert jobs == {} + assert cos == [] + assert fns == [] + assert not path.exists() + + def test_classify_groups_jobs_by_executor(self, tmp_path, monkeypatch): + monkeypatch.setattr(cleaner, 'CLEANER_DIR', str(tmp_path)) + for name, jobs in ( + ('a.pkl', {'ex-0-M000'}), + ('b.pkl', {'ex-0-M001'}), + ('c.pkl', {'other-1-M000'}), + ): + (tmp_path / name).write_bytes(pickle.dumps({ + 'jobs_to_clean': jobs, + 'storage_config': {}, + 'clean_cloudobjects': False, + })) + (tmp_path / 'cos.pkl').write_bytes(pickle.dumps({ + 'cos_to_clean': [], + 'storage_config': {}, + })) + (tmp_path / 'fn.pkl').write_bytes(pickle.dumps({ + 'fn_to_clean': 'ex-0', + 'storage_config': {}, + })) + jobs, cos, fns = _classify_cleaner_files( + ['a.pkl', 'b.pkl', 'c.pkl', 'cos.pkl', 'fn.pkl'] + ) + assert set(jobs) == {'ex-0', 'other-1'} + assert len(jobs['ex-0']) == 2 + assert len(cos) == 1 + assert len(fns) == 1 + + def test_clean_executor_jobs_reuses_storage(self, tmp_path): + files = [] + for name in ('one.pkl', 'two.pkl'): + path = tmp_path / name + path.write_text('x') + files.append({ + 'file_location': str(path), + 'data': { + 'storage_config': {'k': 1}, + 'clean_cloudobjects': True, + 'jobs_to_clean': {'ex-j0'}, + }, + }) + storage = MagicMock() + storage.bucket = 'bkt' + with patch('lithops.scripts.cleaner.Storage', return_value=storage) as st: + with patch('lithops.scripts.cleaner.clean_bucket') as cb: + clean_executor_jobs('ex', files) + st.assert_called_once() + assert cb.call_count == 4 + prefixes = [c.args[2] for c in cb.call_args_list] + assert f'{JOBS_PREFIX}/ex-j0/' in prefixes + assert f'{TEMP_PREFIX}/ex-j0/' in prefixes + assert not (tmp_path / 'one.pkl').exists() + + def test_clean_cloudobjects_same_backend_only(self, tmp_path): + path = tmp_path / 'cos.pkl' + path.write_text('x') + keep = SimpleNamespace(backend='s3', bucket='b', key='keep') + drop = SimpleNamespace(backend='localhost', bucket='b', key='drop') + storage = MagicMock() + storage.backend = 'localhost' + with patch('lithops.scripts.cleaner.Storage', return_value=storage): + clean_cloudobjects({ + 'file_location': str(path), + 'data': { + 'cos_to_clean': [keep, drop], + 'storage_config': {}, + }, + }) + storage.delete_object.assert_called_once_with('b', 'drop') + assert not path.exists() + + def test_clean_functions_deletes_keys(self, tmp_path): + path = tmp_path / 'fn.pkl' + path.write_text('x') + storage = MagicMock() + storage.bucket = 'bkt' + storage.list_keys.return_value = ['k1'] + with patch('lithops.scripts.cleaner.Storage', return_value=storage): + clean_functions({ + 'file_location': str(path), + 'data': { + 'fn_to_clean': 'ex-0', + 'storage_config': {}, + }, + }) + storage.delete_objects.assert_called_once_with('bkt', ['k1']) + + def test_run_clean_tasks_surfaces_exceptions(self): + with patch( + 'lithops.scripts.cleaner.clean_cloudobjects', + side_effect=RuntimeError('boom'), + ): + with pytest.raises(RuntimeError, match='boom'): + _run_clean_tasks( + {}, [{'file_location': 'a', 'data': {}}], [] + ) + + def test_clean_loop_exits_when_idle(self, monkeypatch): + monkeypatch.setattr(cleaner, '_IDLE_CONFIRM_SECONDS', 0) + monkeypatch.setattr(cleaner.os, 'listdir', lambda _: []) + cleaner.clean() + + def test_clean_loop_ignores_log_and_pid_files(self, monkeypatch): + names = [ + os.path.basename(cleaner.CLEANER_LOG_FILE), + os.path.basename(cleaner.CLEANER_PID_FILE), + ] + monkeypatch.setattr(cleaner, '_IDLE_CONFIRM_SECONDS', 0) + monkeypatch.setattr(cleaner.os, 'listdir', lambda _: names) + with patch.object(cleaner, '_run_clean_tasks') as run: + cleaner.clean() + run.assert_not_called() + + def test_clean_loop_picks_up_request_dropped_while_idle(self, monkeypatch): + calls = {'n': 0} + + def listdir(_): + calls['n'] += 1 + if calls['n'] == 2: + return ['late.pkl'] + return [] + + monkeypatch.setattr(cleaner, '_IDLE_CONFIRM_SECONDS', 0) + monkeypatch.setattr(cleaner.os, 'listdir', listdir) + monkeypatch.setattr(cleaner.time, 'sleep', lambda _s: None) + with patch.object( + cleaner, '_classify_cleaner_files', return_value=({}, [], []) + ) as classify: + with patch.object(cleaner, '_run_clean_tasks'): + cleaner.clean() + classify.assert_called_once_with(['late.pkl']) + + def test_classify_discards_unreadable_request(self, tmp_path, monkeypatch): + monkeypatch.setattr(cleaner, 'CLEANER_DIR', str(tmp_path)) + corrupt = tmp_path / 'half-written' + corrupt.write_bytes(b'\x80\x05}') + + assert cleaner._classify_cleaner_files(['half-written']) == ({}, [], []) + assert not corrupt.exists() + + def test_classify_discards_unknown_request(self, tmp_path, monkeypatch): + monkeypatch.setattr(cleaner, 'CLEANER_DIR', str(tmp_path)) + unknown = tmp_path / 'unknown' + with unknown.open('wb') as fh: + pickle.dump({'something_else': 1}, fh) + + assert cleaner._classify_cleaner_files(['unknown']) == ({}, [], []) + assert not unknown.exists() + + def test_pending_requests_ignore_files_being_written( + self, tmp_path, monkeypatch + ): + monkeypatch.setattr(cleaner, 'CLEANER_DIR', str(tmp_path)) + (tmp_path / 'ready').write_text('x') + (tmp_path / f'staging{cleaner.CLEANER_TMP_SUFFIX}').write_text('x') + + assert cleaner._pending_request_files() == ['ready'] + + def test_pending_requests_tolerate_missing_dir(self, tmp_path, monkeypatch): + monkeypatch.setattr(cleaner, 'CLEANER_DIR', str(tmp_path / 'gone')) + assert cleaner._pending_request_files() == [] + + def test_main_skips_while_another_cleaner_holds_the_lock( + self, tmp_path, monkeypatch + ): + monkeypatch.setattr(cleaner, 'CLEANER_PID_FILE', str(tmp_path / 'cleaner.pid')) + monkeypatch.setattr(cleaner, 'CLEANER_DIR', str(tmp_path)) + monkeypatch.setattr(cleaner, '_LOCK_RETRY_SECONDS', 0) + + held = cleaner._lock_pid_file() + assert held is not None + try: + with patch.object(cleaner, '_configure_cleaner_logging') as cfg: + cleaner.main() + cfg.assert_not_called() + finally: + os.close(held) + + def test_main_takes_the_lock_a_dead_cleaner_left_behind( + self, tmp_path, monkeypatch + ): + pid = tmp_path / 'cleaner.pid' + # The pid of a cleaner that died without removing its file. Nothing + # holds the lock any more, so this run must not be blocked by it + pid.write_text('999999999') + monkeypatch.setattr(cleaner, 'CLEANER_PID_FILE', str(pid)) + monkeypatch.setattr(cleaner, 'CLEANER_DIR', str(tmp_path)) + + with patch.object(cleaner, '_configure_cleaner_logging'): + with patch.object(cleaner, 'clean') as clean_fn: + cleaner.main() + + clean_fn.assert_called_once() + assert pid.read_text() == str(os.getpid()) + + def test_lock_is_released_when_the_cleaner_finishes( + self, tmp_path, monkeypatch + ): + monkeypatch.setattr(cleaner, 'CLEANER_PID_FILE', str(tmp_path / 'cleaner.pid')) + monkeypatch.setattr(cleaner, 'CLEANER_DIR', str(tmp_path)) + + with patch.object(cleaner, '_configure_cleaner_logging'): + with patch.object(cleaner, 'clean'): + cleaner.main() + + second = cleaner._lock_pid_file() + assert second is not None + os.close(second) diff --git a/lithops/tests/test_serverless.py b/lithops/tests/test_serverless.py new file mode 100644 index 000000000..3a97111ec --- /dev/null +++ b/lithops/tests/test_serverless.py @@ -0,0 +1,122 @@ +# +# Unit tests for the serverless compute frontend (not cloud backends). +# + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from lithops.serverless import ServerlessHandler +from lithops.serverless import __all__ as serverless_all +from lithops.utils import BackendType + + +def _config(**extra): + cfg = { + 'backend': 'aws_lambda', + 'aws_lambda': {'region': 'us-east-1'}, + } + cfg.update(extra) + return cfg + + +def _make_handler(backend=None): + fake_backend = backend or MagicMock() + fake_backend.type = BackendType.FAAS.value + fake_module = MagicMock() + fake_module.ServerlessBackend.return_value = fake_backend + storage = MagicMock() + with patch( + 'lithops.serverless.serverless.importlib.import_module', + return_value=fake_module, + ) as importer: + handler = ServerlessHandler(_config(), storage) + fake_module.ServerlessBackend.assert_called_once_with( + {'region': 'us-east-1'}, storage + ) + importer.assert_called_once_with('lithops.serverless.backends.aws_lambda') + handler.backend = fake_backend + return handler, fake_backend + + +class TestServerlessExports: + + def test_all_exports_handler(self): + assert serverless_all == ['ServerlessHandler'] + + +class TestServerlessHandler: + + def test_loads_named_backend_and_reports_type(self): + handler, backend = _make_handler() + assert handler.backend_name == 'aws_lambda' + assert handler.get_backend_type() == BackendType.FAAS.value + handler.init() + backend.init.assert_not_called() + + def test_backend_import_failure_is_logged_and_reraised(self): + with patch( + 'lithops.serverless.serverless.importlib.import_module', + side_effect=ImportError('missing extra'), + ): + with pytest.raises(ImportError, match='missing extra'): + ServerlessHandler(_config(), MagicMock()) + + def test_invoke_passes_runtime_fields(self): + handler, backend = _make_handler() + payload = { + 'runtime_name': 'lithops/python:3.12', + 'runtime_memory': 256, + 'call_ids': ['00000'], + } + backend.invoke.return_value = {'ok': True} + assert handler.invoke(payload) == {'ok': True} + backend.invoke.assert_called_once_with( + 'lithops/python:3.12', 256, payload + ) + + def test_pre_invoke_and_clear_are_optional(self): + class MinimalBackend: + type = BackendType.BATCH.value + + handler, _ = _make_handler(MinimalBackend()) + handler.pre_invoke( + SimpleNamespace(runtime_name='rt', runtime_memory=128) + ) + handler.clear(['job-1']) + + def test_pre_invoke_and_clear_delegate_when_present(self): + handler, backend = _make_handler() + job = SimpleNamespace(runtime_name='rt', runtime_memory=128) + handler.pre_invoke(job) + backend.pre_invoke.assert_called_once_with('rt', 128) + handler.clear(['job-1'], exception=RuntimeError('x')) + backend.clear.assert_called_once_with(['job-1']) + + def test_build_runtime_defaults_extra_args_to_empty_list(self): + handler, backend = _make_handler() + handler.build_runtime('rt', None) + backend.build_runtime.assert_called_once_with('rt', None, []) + + def test_runtime_lifecycle_delegates(self): + handler, backend = _make_handler() + backend.deploy_runtime.return_value = {'preinstalls': []} + backend.list_runtimes.return_value = ['rt'] + backend.get_runtime_key.return_value = 'key' + backend.get_runtime_info.return_value = {'runtime_name': 'rt'} + + assert handler.deploy_runtime('rt', 256, 60) == {'preinstalls': []} + backend.deploy_runtime.assert_called_once_with('rt', 256, timeout=60) + + handler.delete_runtime('rt', 256, '3.0.0') + backend.delete_runtime.assert_called_once_with('rt', 256, '3.0.0') + + handler.clean(all=True) + backend.clean.assert_called_once_with(all=True) + + assert handler.list_runtimes() == ['rt'] + backend.list_runtimes.assert_called_once_with('all') + + assert handler.get_runtime_key('rt', 256, '3.0.0') == 'key' + assert handler.get_runtime_info() == {'runtime_name': 'rt'} diff --git a/lithops/tests/test_standalone.py b/lithops/tests/test_standalone.py new file mode 100644 index 000000000..78bc0207e --- /dev/null +++ b/lithops/tests/test_standalone.py @@ -0,0 +1,664 @@ +# +# Unit tests for the standalone compute frontend (not cloud backends). +# + +import json +import os +import sys +from unittest.mock import MagicMock, patch + +import pytest + +from lithops.standalone import LithopsValidationError, StandaloneHandler +from lithops.standalone import __all__ as standalone_all +from lithops.standalone.keeper import BudgetKeeper +from lithops.standalone.runner import main as runner_main +from lithops.standalone.standalone import StandaloneHandler as HandlerCls +from lithops.standalone.utils import ( + docker_login, + JobStatus, + LithopsValidationError as UtilsError, + StandaloneMode, + _format_apt_packages_for_shell, + get_host_setup_script, + get_master_setup_script, + get_worker_setup_script, + is_container_runtime, + lithops_pip_spec_from_config, +) +from lithops.utils import BackendType +from lithops.version import __version__ + + +def _handler_config(exec_mode='consume', **extra): + cfg = { + 'backend': 'vm', + 'exec_mode': exec_mode, + 'start_timeout': 5, + 'runtime': 'python3', + 'hard_dismantle_timeout': 60, + 'vm': {'max_workers': 2}, + } + cfg.update(extra) + return cfg + + +def _make_handler(exec_mode='consume', **extra): + fake_backend = MagicMock() + fake_module = MagicMock() + fake_module.StandaloneBackend.return_value = fake_backend + with patch( + 'lithops.standalone.standalone.importlib.import_module', + return_value=fake_module, + ): + handler = StandaloneHandler(_handler_config(exec_mode, **extra)) + handler.backend = fake_backend + return handler + + +def _keeper_config(**extra): + cfg = { + 'auto_dismantle': True, + 'soft_dismantle_timeout': 10, + 'hard_dismantle_timeout': 20, + 'exec_mode': 'reuse', + } + cfg.update(extra) + return cfg + + +def _make_keeper(**keeper_kwargs): + instance = MagicMock() + instance.name = 'worker-1' + instance.private_ip = '10.0.0.8' + instance.instance_id = 'i-1' + instance.delete_on_dismantle = False + with patch('lithops.standalone.keeper.StandaloneHandler') as handler_cls: + handler_cls.return_value.backend.get_instance.return_value = instance + keeper = BudgetKeeper(_keeper_config(), {'name': 'worker-1'}, **keeper_kwargs) + keeper.instance = instance + return keeper, instance + + +class TestStandaloneExports: + + def test_all_is_string_names(self): + assert standalone_all == ['StandaloneHandler', 'LithopsValidationError'] + assert LithopsValidationError is UtilsError + assert HandlerCls is StandaloneHandler + + +class TestStandaloneUtils: + + def test_pip_spec_defaults_and_cloud_extras(self): + assert lithops_pip_spec_from_config(None) == 'lithops' + assert lithops_pip_spec_from_config({}) == 'lithops' + assert lithops_pip_spec_from_config({'backend': 'localhost'}) == ( + 'lithops[redis]' + ) + spec = lithops_pip_spec_from_config({'lithops': {'backend': 'aws_ec2'}}) + assert spec == 'lithops[aws,redis]' + + def test_apt_packages_reject_invalid_names(self): + with pytest.raises(LithopsValidationError, match='apt package'): + _format_apt_packages_for_shell('foo; rm -rf /') + + def test_docker_login_is_empty_without_credentials(self): + assert docker_login({'backend': 'vm', 'vm': {}}) == '' + assert docker_login( + {'backend': 'vm', 'vm': {'docker_user': 'me'}} + ) == '' + + def test_docker_login_keeps_the_password_out_of_the_process_list(self): + script = docker_login({'backend': 'vm', 'vm': { + 'docker_server': 'reg.io', + 'docker_user': 'me', + 'docker_password': 'secret', + }}) + # -p would show the password to anyone running ps on the VM + assert '--password-stdin' in script + assert '-p secret' not in script + assert 'docker login -u me' in script + assert '/opt/lithops/setup.log' in script + + def test_docker_login_quotes_the_credentials_it_is_given(self): + # Unquoted, the ; would run "rd" as a command on the VM + script = docker_login({'backend': 'vm', 'vm': { + 'docker_server': 'reg.io', + 'docker_user': 'me', + 'docker_password': 'p@ss w;rd', + }}) + assert "'p@ss w;rd'" in script + assert '; rd' not in script.replace("'p@ss w;rd'", '') + + def test_container_runtime_detects_docker_tags(self): + assert is_container_runtime('python3') is False + assert is_container_runtime('/usr/bin/python3') is False + assert is_container_runtime('python:3.12') is True + assert is_container_runtime('lithops/python:3.12') is True + + def test_worker_setup_script_uses_native_python(self): + script = get_worker_setup_script( + {'backend': 'vm', 'runtime': 'python3', 'use_gpu': False, 'vm': {}}, + {'master_ip': '10.0.0.1'}, + ) + assert '/usr/bin/python3' in script + assert 'docker run --rm --name lithops_worker' not in script + + def test_worker_setup_script_uses_docker_for_tagged_python(self): + script = get_worker_setup_script( + { + 'backend': 'vm', + 'runtime': 'python:3.12', + 'use_gpu': True, + 'vm': {}, + }, + {'master_ip': '10.0.0.1'}, + ) + assert 'docker run --rm --name lithops_worker' in script + assert '--gpus all' in script + assert 'python:3.12' in script + assert '-v /opt/lithops:/opt/lithops' in script + assert '-v /tmp:/tmp' in script + assert '/opt/lithops/worker.py' in script + assert '\\opt\\lithops' not in script + assert '\\tmp\\' not in script + + def test_host_and_master_setup_scripts_use_posix_remote_paths(self): + host = get_host_setup_script() + assert host.startswith('#!/bin/bash') + assert 'mkdir -p /opt/lithops' in host + assert '/opt/lithops/setup.log' in host + assert '/opt/lithops/setup-done.flag' in host + assert '\\opt\\lithops' not in host + + master = get_master_setup_script( + {'backend': 'vm', 'vm': {}}, + {'master_ip': '10.0.0.1'}, + ) + assert 'unzip -o /tmp/lithops_standalone.zip -d /opt/lithops' in master + assert '/opt/lithops/master.data' in master + assert '/opt/lithops/config' in master + assert '\\opt\\lithops' not in master + assert '\\tmp\\lithops_standalone.zip' not in master + + +class TestBudgetKeeper: + + def test_running_attribute_and_job_tracking(self): + keeper, _instance = _make_keeper() + assert keeper.running is False + assert not hasattr(keeper, 'runing') + keeper.add_job('job-a') + assert keeper.jobs['job-a'] == JobStatus.RUNNING.value + keeper.set_job_done('job-a') + assert keeper.jobs['job-a'] == JobStatus.DONE.value + assert keeper._all_jobs_done() is True + + def test_mark_finished_jobs_survives_a_concurrent_add(self): + # The service adds jobs from its own threads, so iterating the dict + # itself would raise "dictionary changed size during iteration" and + # kill the keeper, leaving the instance running forever + keeper, _instance = _make_keeper() + keeper.add_job('job-a') + + real_isfile = os.path.isfile + + def isfile(path): + keeper.jobs[f'job-{len(keeper.jobs)}'] = JobStatus.RUNNING.value + return real_isfile(path) + + with patch('lithops.standalone.keeper.os.path.isfile', side_effect=isfile): + keeper._mark_finished_jobs() + assert len(keeper.jobs) > 1 + + def test_stop_instance_calls_stop_callback(self): + stop = MagicMock() + delete = MagicMock() + keeper, instance = _make_keeper(stop_callback=stop, delete_callback=delete) + instance.delete_on_dismantle = False + keeper.stop_instance() + stop.assert_called_once() + delete.assert_not_called() + instance.stop.assert_called_once() + assert keeper.running is False + + def test_stop_instance_calls_delete_callback(self): + stop = MagicMock() + delete = MagicMock() + keeper, instance = _make_keeper(stop_callback=stop, delete_callback=delete) + instance.delete_on_dismantle = True + keeper.stop_instance() + delete.assert_called_once() + stop.assert_not_called() + + +class TestStandaloneHandler: + + def test_backend_type_and_runtime_info(self): + handler = _make_handler() + assert handler.get_backend_type() == BackendType.BATCH.value + assert handler.exec_mode is StandaloneMode.CONSUME + assert handler.get_runtime_info() == { + 'runtime_name': 'python3', + 'runtime_memory': None, + 'runtime_timeout': 60, + 'max_workers': 2, + } + + def test_build_image_defaults_extra_args_to_empty_list(self): + handler = _make_handler() + handler.build_image('img', None, False, None) + handler.backend.build_image.assert_called_once_with( + 'img', None, False, None, [] + ) + + def test_get_runtime_key_delegates_to_backend(self): + handler = _make_handler() + handler.backend.get_runtime_key.return_value = 'key' + assert handler.get_runtime_key('python3', None) == 'key' + handler.backend.get_runtime_key.assert_called_once_with( + 'python3', __version__ + ) + + def test_master_ready_rejects_version_mismatch(self): + handler = _make_handler() + handler._make_request = MagicMock(return_value={'response': '0.0.0'}) + with pytest.raises(LithopsValidationError, match='doesn\'t match'): + handler._is_master_service_ready() + + def test_create_workers_zero_is_noop(self): + handler = _make_handler('create') + assert handler._create_workers(0, 'e', 'M000') == [] + handler.backend.create_worker.assert_not_called() + + def test_invoke_consume_uses_master_as_worker(self): + handler = _make_handler('consume') + master = MagicMock() + master.name = 'master' + master.private_ip = '10.0.0.2' + master.instance_id = 'i-m' + master.ssh_credentials = {'username': 'ubuntu'} + master.instance_type = 'unused' + handler.backend.master = master + handler._is_master_service_ready = MagicMock(return_value=True) + handler._make_request = MagicMock() + payload = { + 'executor_id': 'sess-0', + 'job_id': 'M000', + 'job_key': 'sess-0-M000', + 'total_calls': 2, + 'worker_processes': 1, + 'config': {'lithops': {'backend': 'vm'}, 'vm': {'ssh_key_filename': 'k'}}, + } + handler.invoke(payload) + assert payload['worker_instances'] == [{ + 'name': 'master', + 'private_ip': '10.0.0.2', + 'instance_id': 'i-m', + 'ssh_credentials': {'username': 'ubuntu'}, + 'instance_type': 'unused', + }] + assert 'ssh_key_filename' not in payload['config']['vm'] + handler._make_request.assert_called_once_with('POST', 'job/run', payload) + assert handler.jobs == ['sess-0-M000'] + + def _create_payload(self, **extra): + payload = { + 'executor_id': 'sess-0', + 'job_id': 'M000', + 'job_key': 'sess-0-M000', + 'total_calls': 5, + 'worker_processes': 2, + 'max_workers': 10, + 'runtime_name': 'python3', + 'config': {'lithops': {'backend': 'vm'}, 'vm': {}}, + } + payload.update(extra) + return payload + + def _ready_handler(self, mode): + handler = _make_handler(mode) + handler.backend.get_worker_instance_type.return_value = 'big' + handler.backend.get_worker_cpu_count.return_value = 2 + handler._is_master_service_ready = MagicMock(return_value=True) + handler._make_request = MagicMock(return_value=[]) + return handler + + def test_invoke_create_rounds_workers_up(self): + handler = self._ready_handler('create') + created = [] + + def create_workers(count, executor_id, job_id): + created.append((count, executor_id, job_id)) + return [MagicMock(name=f'w{n}') for n in range(count)] + + handler._create_workers = MagicMock(side_effect=create_workers) + payload = self._create_payload() + handler.invoke(payload) + # 5 calls over 2 processes per worker needs 3 workers + assert created == [(3, 'sess-0', 'M000')] + assert payload['worker_instance_type'] == 'big' + assert len(payload['worker_instances']) == 3 + + def test_invoke_create_caps_workers_at_max_workers(self): + handler = self._ready_handler('create') + handler._create_workers = MagicMock(return_value=[MagicMock()]) + handler.invoke(self._create_payload(total_calls=100, max_workers=4)) + assert handler._create_workers.call_args[0][0] == 4 + + def test_invoke_create_resolves_auto_worker_processes(self): + handler = self._ready_handler('create') + handler._create_workers = MagicMock(return_value=[MagicMock()]) + payload = self._create_payload(worker_processes='AUTO') + handler.invoke(payload) + assert payload['worker_processes'] == 2 + assert payload['config']['vm']['worker_processes'] == 2 + + def test_invoke_reuse_only_creates_the_missing_workers(self): + handler = self._ready_handler('reuse') + handler._get_workers_on_master = MagicMock(return_value=['w1']) + handler._create_workers = MagicMock(return_value=[MagicMock()]) + handler.invoke(self._create_payload()) + # 3 needed, 1 already free on the master + assert handler._create_workers.call_args[0][0] == 2 + + def test_invoke_reuse_creates_nothing_when_enough_workers(self): + handler = self._ready_handler('reuse') + handler._get_workers_on_master = MagicMock( + return_value=['w1', 'w2', 'w3'] + ) + handler._create_workers = MagicMock() + payload = self._create_payload() + handler.invoke(payload) + handler._create_workers.assert_not_called() + assert payload['worker_instances'] == [] + + def test_invoke_raises_when_no_worker_could_be_created(self): + handler = self._ready_handler('create') + handler._create_workers = MagicMock(return_value=[]) + with pytest.raises(Exception, match='not possible to create any workers'): + handler.invoke(self._create_payload()) + + def test_invoke_sets_up_the_master_when_it_is_not_ready(self): + handler = self._ready_handler('consume') + handler._is_master_service_ready = MagicMock(return_value=False) + handler._validate_master_service_setup = MagicMock() + handler._wait_master_service_ready = MagicMock() + handler.invoke(self._create_payload(worker_processes=1)) + handler.backend.master.create.assert_called_once_with(check_if_exists=True) + handler.backend.master.wait_ready.assert_called_once() + handler._validate_master_service_setup.assert_called_once() + handler._wait_master_service_ready.assert_called_once() + + def test_request_from_worker_uses_lithops_master_host(self): + handler = _make_handler() + handler.is_lithops_worker = True + with patch('lithops.standalone.standalone.requests.get') as get: + get.return_value.json.return_value = {'response': __version__} + assert handler._make_request('GET', 'ping') == { + 'response': __version__ + } + assert 'lithops-master' in get.call_args[0][0] + + def test_request_from_worker_accepts_an_empty_post_body(self): + handler = _make_handler() + handler.is_lithops_worker = True + with patch('lithops.standalone.standalone.requests.post') as post: + post.return_value.content = b'' + assert handler._make_request( + 'POST', 'job/stop', ['sess-0-M000'] + ) is None + post.return_value.raise_for_status.assert_called_once() + + def test_request_via_ssh_accepts_an_empty_response(self): + # /job/stop and /clean answer 204 with no body. curl used to print + # its progress meter on stderr, which was raised as a failure + handler = _make_handler() + ssh = handler.backend.master.get_ssh_client.return_value + ssh.run_remote_command.return_value = ('', '') + assert handler._make_request( + 'POST', 'job/stop', ['sess-0-M000'] + ) is None + cmd = ssh.run_remote_command.call_args[0][0] + assert cmd.startswith('curl -sS ') + assert 'job/stop' in cmd + + def test_request_via_ssh_raises_when_curl_prints_an_error(self): + handler = _make_handler() + ssh = handler.backend.master.get_ssh_client.return_value + ssh.run_remote_command.return_value = ( + '', 'curl: (7) Failed to connect' + ) + with pytest.raises(ValueError, match='Failed to connect'): + handler._make_request('POST', 'job/stop', ['sess-0-M000']) + + def test_request_via_ssh_parses_a_json_body(self): + handler = _make_handler() + ssh = handler.backend.master.get_ssh_client.return_value + ssh.run_remote_command.return_value = ('{"response": "ok"}', '') + assert handler._make_request('GET', 'ping') == {'response': 'ok'} + + def test_clear_logs_that_jobs_were_stopped(self): + handler = _make_handler('reuse') + handler.jobs = ['sess-0-M000'] + handler._make_request = MagicMock(return_value=None) + with patch('lithops.standalone.standalone.logger') as log: + handler.clear() + handler._make_request.assert_called_once_with( + 'POST', 'job/stop', ['sess-0-M000'] + ) + log.debug.assert_called_once_with('Jobs stopped on the master') + handler.backend.clear.assert_not_called() + + def test_clear_logs_when_the_master_cannot_be_reached(self): + handler = _make_handler('reuse') + handler._make_request = MagicMock(side_effect=ValueError('down')) + with patch('lithops.standalone.standalone.logger') as log: + handler.clear() + log.debug.assert_called_once() + assert log.debug.call_args[0][0].startswith( + 'Could not stop the jobs on the master:' + ) + + +class TestStandaloneRunner: + + def test_import_does_not_open_log_stream(self): + from lithops.standalone import runner as sa_runner + assert getattr(sa_runner, 'log_file_stream', None) is None + + def test_run_job_reads_text_json(self, tmp_path, monkeypatch): + from lithops.standalone import runner as sa_runner + monkeypatch.setattr(sa_runner, 'RN_LOG_FILE', str(tmp_path / 'rn.log')) + task = tmp_path / 'task.json' + task.write_text(json.dumps({ + 'executor_id': 'sess-0', + 'job_id': 'M000', + 'call_ids': ['00000'], + })) + monkeypatch.setattr(sys, 'argv', ['runner.py', 'aws_ec2', str(task)]) + monkeypatch.setenv('__LITHOPS_BACKEND', '') + monkeypatch.setenv('__LITHOPS_ACTIVATION_ID', '') + with patch.object(sa_runner, 'function_handler') as handler: + runner_main() + payload = handler.call_args[0][0] + assert payload['worker_processes'] == 1 + assert os.environ['__LITHOPS_BACKEND'] == 'AWS EC2' + + +class TestStandaloneMasterWorkerHttp: + """These tests need flask, which is not required for localhost runs.""" + + @pytest.fixture(autouse=True) + def _need_flask(self): + pytest.importorskip('flask') + + def test_master_rejects_non_dict_metadata(self): + from lithops.standalone import master as sa_master + sa_master.budget_keeper = MagicMock() + client = sa_master.app.test_client() + resp = client.get('/metadata', json=['not-a-dict']) + assert resp.status_code == 404 + assert 'dictionary' in resp.get_json()['error'] + + def test_master_map_if_any_reports_failures_without_raising(self): + from lithops.standalone import master as sa_master + + def boom(item): + raise RuntimeError(f'no {item}') + + with patch.object(sa_master, 'logger') as log: + sa_master._map_if_any(boom, ['worker:a', 'worker:b']) + assert log.error.call_count == 2 + assert 'no worker:a' in log.error.call_args_list[0][0][0] + + def test_master_cancel_job_survives_an_emptied_queue(self): + from lithops.standalone import master as sa_master + redis_client = MagicMock() + redis_client.hget.return_value = 'wq:sess-0-M000' + # A worker took the last task between llen() and rpop() + redis_client.llen.return_value = 1 + redis_client.rpop.return_value = None + redis_client.keys.return_value = [] + with patch.object(sa_master, 'redis_client', redis_client): + sa_master.cancel_job_process(['sess-0-M000']) + redis_client.rpop.assert_called_once() + + def test_master_cancel_job_skips_a_job_with_no_queue(self): + from lithops.standalone import master as sa_master + redis_client = MagicMock() + redis_client.hget.return_value = None + with patch.object(sa_master, 'redis_client', redis_client): + sa_master.cancel_job_process(['sess-0-M000']) + redis_client.llen.assert_not_called() + + def test_worker_ping_counts_idle_and_busy(self): + from lithops.standalone import worker as sa_worker + from lithops.standalone.utils import WorkerStatus + sa_worker.worker_threads = { + 0: {'status': WorkerStatus.IDLE.value}, + 1: {'status': WorkerStatus.BUSY.value}, + 2: {'status': WorkerStatus.IDLE.value}, + } + client = sa_worker.app.test_client() + assert client.get('/ping').get_json() == {'busy': 1, 'free': 2} + + def test_worker_stop_survives_a_concurrent_task_registration(self): + from lithops.standalone import worker as sa_worker + # The consumer threads add and remove entries while this route runs, + # so iterating the dict itself raised "dictionary changed size during + # iteration" and failed the stop request + proc = MagicMock() + proc.pid = 7 + sa_worker.job_processes = {'sess-0-M000-00000': proc} + sa_worker.canceled = [] + + def kill(process): + sa_worker.job_processes['sess-0-M000-00001'] = MagicMock() + + client = sa_worker.app.test_client() + with patch.object(sa_worker, '_kill_process_group', side_effect=kill), \ + patch.object(sa_worker.Path, 'touch'): + resp = client.post('/stop/sess-0-M000') + + assert resp.status_code == 200 + assert 'sess-0-M000' in sa_worker.canceled + assert 'sess-0-M000-00000' not in sa_worker.job_processes + + def test_worker_ttd_disabled_without_keeper(self): + from lithops.standalone import worker as sa_worker + sa_worker.budget_keeper = None + client = sa_worker.app.test_client() + resp = client.get('/ttd') + assert resp.status_code == 200 + assert resp.get_data(as_text=True) == 'Disabled' + + def test_wait_for_task_in_reuse_mode_polls_with_a_timeout(self): + from lithops.standalone import worker as sa_worker + redis_client = MagicMock() + redis_client.brpop.return_value = None + with patch.object(sa_worker, 'redis_client', redis_client): + assert sa_worker._wait_for_task( + 'wq:t3.micro-2-python3', StandaloneMode.REUSE.value + ) is None + redis_client.brpop.assert_called_once_with( + 'wq:t3.micro-2-python3', timeout=sa_worker._QUEUE_POLL_TIMEOUT + ) + redis_client.rpop.assert_not_called() + + def test_wait_for_task_in_reuse_mode_returns_the_payload(self): + from lithops.standalone import worker as sa_worker + redis_client = MagicMock() + redis_client.brpop.return_value = ('wq', '{"call_ids": ["00000"]}') + with patch.object(sa_worker, 'redis_client', redis_client): + assert sa_worker._wait_for_task( + 'wq', StandaloneMode.CONSUME.value + ) == '{"call_ids": ["00000"]}' + + def test_wait_for_task_in_create_mode_stops_on_an_empty_queue(self): + from lithops.standalone import worker as sa_worker + redis_client = MagicMock() + redis_client.rpop.return_value = None + with patch.object(sa_worker, 'redis_client', redis_client): + assert sa_worker._wait_for_task( + 'wq', StandaloneMode.CREATE.value + ) is None + redis_client.rpop.assert_called_once_with('wq') + redis_client.brpop.assert_not_called() + + def test_consumer_stays_idle_across_empty_reuse_polls(self): + from lithops.standalone import worker as sa_worker + from lithops.standalone.utils import WorkerStatus + + class StopConsumer(BaseException): + pass + + sa_worker.worker_threads = {0: {'status': None}} + polls = {'n': 0} + + def wait(queue_name, exec_mode): + polls['n'] += 1 + assert sa_worker.worker_threads[0]['status'] == WorkerStatus.IDLE.value + if polls['n'] < 3: + return None + raise StopConsumer + + with patch.object(sa_worker, '_wait_for_task', side_effect=wait): + try: + sa_worker.redis_queue_consumer( + 0, 'wq', StandaloneMode.REUSE.value, 'aws_ec2' + ) + except StopConsumer: + pass + + assert polls['n'] == 3 + assert sa_worker.worker_threads[0]['status'] == WorkerStatus.IDLE.value + + def test_consumer_retries_after_a_lost_redis_connection(self): + from lithops.standalone import worker as sa_worker + from lithops.standalone.utils import WorkerStatus + + class StopConsumer(BaseException): + pass + + sa_worker.worker_threads = {0: {'status': None}} + polls = {'n': 0} + + def wait(queue_name, exec_mode): + polls['n'] += 1 + if polls['n'] == 1: + raise ConnectionError('timed out') + raise StopConsumer + + with patch.object(sa_worker, '_wait_for_task', side_effect=wait), \ + patch.object(sa_worker.time, 'sleep'): + try: + sa_worker.redis_queue_consumer( + 0, 'wq', StandaloneMode.REUSE.value, 'aws_ec2' + ) + except StopConsumer: + pass + + assert polls['n'] == 2 + assert sa_worker.worker_threads[0]['status'] == WorkerStatus.IDLE.value diff --git a/lithops/tests/test_standalone_master.py b/lithops/tests/test_standalone_master.py new file mode 100644 index 000000000..701fab8cc --- /dev/null +++ b/lithops/tests/test_standalone_master.py @@ -0,0 +1,824 @@ +# +# Unit tests for the standalone master service (not cloud backends). +# +# The master runs as a Flask service on the master VM, keeps its state in +# redis, and reaches the workers over HTTP, so every test here drives it with +# those three replaced. +# + +import json +from contextlib import ExitStack +from unittest.mock import MagicMock, patch + +import pytest + +from lithops.standalone.utils import JobStatus, StandaloneMode, WorkerStatus +from lithops.version import __version__ + +pytest.importorskip('flask') +pytest.importorskip('gevent') +pytest.importorskip('redis') + +from lithops.standalone import master as sa_master # noqa: E402 + + +@pytest.fixture(autouse=True) +def master_globals(): + """ + Gives the module its globals back after each test: they are module level + and the service sets them once, in main() + """ + saved = ( + sa_master.redis_client, + sa_master.budget_keeper, + sa_master.master_ip, + ) + sa_master.redis_client = MagicMock() + sa_master.budget_keeper = MagicMock() + sa_master.master_ip = '10.0.0.1' + yield + ( + sa_master.redis_client, + sa_master.budget_keeper, + sa_master.master_ip, + ) = saved + + +@pytest.fixture +def client(): + return sa_master.app.test_client() + + +def _worker_instance(name='lithops-worker-1', **extra): + worker = MagicMock() + worker.name = name + worker.private_ip = extra.get('private_ip', '10.0.0.2') + worker.instance_id = extra.get('instance_id', 'i-1') + worker.instance_type = extra.get('instance_type', 'big') + worker.ssh_credentials = {'username': 'ubuntu'} + worker.config = extra.get('config', {'worker_processes': 2}) + return worker + + +def _standalone_config(**extra): + cfg = { + 'backend': 'vm', + 'exec_mode': StandaloneMode.CREATE.value, + 'runtime': 'python3', + 'vm': {'secret': 'do-not-store'}, + } + cfg.update(extra) + return cfg + + +class TestMasterEndpoints: + + def test_ping_answers_with_the_lithops_version(self, client): + resp = client.get('/ping') + assert resp.status_code == 200 + assert resp.get_json() == {'response': __version__} + + def test_error_reports_the_message(self): + with sa_master.app.app_context(): + resp = sa_master.error('nope') + assert resp.status_code == 404 + assert resp.get_json() == {'error': 'nope'} + + def test_clean_drops_everything_in_redis(self, client): + resp = client.post('/clean') + assert resp.status_code == 204 + sa_master.redis_client.flushall.assert_called_once() + + def test_metadata_rejects_a_body_that_is_not_a_dict(self, client): + resp = client.get('/metadata', json=['not-a-dict']) + assert resp.status_code == 404 + assert 'dictionary' in resp.get_json()['error'] + + def test_metadata_rejects_an_invalid_runtime_name(self, client): + # A space cannot appear in a container image name + resp = client.get('/metadata', json={'runtime': 'has space'}) + assert resp.status_code == 404 + assert 'not valid' in resp.get_json()['error'] + + def test_metadata_returns_what_the_runtime_reports(self, client): + handler = MagicMock() + handler.deploy_runtime.return_value = { + 'lithops_version': __version__, 'preinstalls': [] + } + with patch.object( + sa_master, 'LocalhostHandler', return_value=handler + ): + resp = client.get('/metadata', json={'runtime': 'python3'}) + assert resp.status_code == 200 + assert resp.get_json()['lithops_version'] == __version__ + handler.init.assert_called_once() + handler.deploy_runtime.assert_called_once_with('python3') + + def test_job_stop_rejects_a_body_that_is_not_a_list(self, client): + resp = client.post('/job/stop', json={'not': 'a list'}) + assert resp.status_code == 404 + assert 'list' in resp.get_json()['error'] + + def test_job_stop_cancels_in_the_background(self, client): + with patch.object(sa_master, 'Thread') as thread: + resp = client.post('/job/stop', json=['sess-0-M000']) + assert resp.status_code == 204 + thread.assert_called_once() + assert thread.call_args.kwargs['args'] == (['sess-0-M000'],) + + +class TestWorkerReachability: + + def test_is_worker_free_when_a_process_is_free(self): + with patch.object(sa_master.requests, 'get') as get: + get.return_value.json.return_value = {'free': 2, 'busy': 0} + assert sa_master.is_worker_free('10.0.0.2') is True + + def test_is_worker_free_is_false_when_every_process_is_busy(self): + with patch.object(sa_master.requests, 'get') as get: + get.return_value.json.return_value = {'free': 0, 'busy': 2} + assert sa_master.is_worker_free('10.0.0.2') is False + + def test_is_worker_free_is_false_when_it_cannot_be_reached(self): + with patch.object( + sa_master.requests, 'get', side_effect=OSError('down') + ): + assert sa_master.is_worker_free('10.0.0.2') is False + + def test_worker_ttd_of_the_master_comes_from_its_own_keeper(self): + sa_master.budget_keeper.get_time_to_dismantle.return_value = 42 + assert sa_master.get_worker_ttd('10.0.0.1') == '42' + + def test_worker_ttd_of_another_worker_is_asked_over_http(self): + with patch.object(sa_master.requests, 'get') as get: + get.return_value.text = '77' + assert sa_master.get_worker_ttd('10.0.0.2') == '77' + assert '10.0.0.2' in get.call_args[0][0] + + def test_worker_ttd_is_unknown_when_it_cannot_be_asked(self): + with patch.object( + sa_master.requests, 'get', side_effect=OSError('down') + ): + assert sa_master.get_worker_ttd('10.0.0.2') == 'Unknown' + + +class TestWorkerListing: + + def _worker_data(self, **extra): + data = { + 'name': 'lithops-worker-1', + 'created': '1700000000', + 'instance_type': 'big', + 'worker_processes': '2', + 'runtime': 'python3', + 'exec_mode': StandaloneMode.REUSE.value, + 'status': WorkerStatus.IDLE.value, + 'private_ip': '10.0.0.2', + 'instance_id': 'i-1', + 'ssh_credentials': '{}', + } + data.update(extra) + return data + + def test_worker_list_builds_a_table_with_a_header(self, client): + sa_master.redis_client.keys.return_value = ['worker:lithops-worker-1'] + sa_master.redis_client.hgetall.return_value = self._worker_data() + with patch.object(sa_master, 'get_worker_ttd', return_value='30'): + resp = client.get('/worker/list') + table = resp.get_json() + assert table[0][0] == 'Worker Name' + assert len(table) == 2 + row = table[1] + assert row[0] == 'lithops-worker-1' + assert row[2] == 'big' + assert row[-1] == '30s' + + def test_worker_list_leaves_a_ttd_word_unsuffixed(self, client): + sa_master.redis_client.keys.return_value = ['worker:lithops-worker-1'] + sa_master.redis_client.hgetall.return_value = self._worker_data() + with patch.object(sa_master, 'get_worker_ttd', return_value='Disabled'): + resp = client.get('/worker/list') + assert resp.get_json()[1][-1] == 'Disabled' + + def test_worker_list_is_just_a_header_with_no_workers(self, client): + sa_master.redis_client.keys.return_value = [] + resp = client.get('/worker/list') + assert len(resp.get_json()) == 1 + + def test_worker_get_rejects_a_body_that_is_not_a_dict(self, client): + sa_master.redis_client.keys.return_value = [] + resp = client.get('/worker/get', json=['nope']) + assert resp.status_code == 404 + + def test_worker_get_returns_only_the_free_workers_of_that_shape( + self, client + ): + wanted = self._worker_data(name='wanted') + other_type = self._worker_data(name='other', instance_type='small') + other_rt = self._worker_data(name='other-rt', runtime='python:3.12') + busy = self._worker_data(name='busy', private_ip='10.0.0.9') + sa_master.redis_client.keys.return_value = ['w1', 'w2', 'w3', 'w4'] + sa_master.redis_client.hgetall.side_effect = [ + wanted, other_type, other_rt, busy + ] + + def free(private_ip): + return private_ip != '10.0.0.9' + + with patch.object(sa_master, 'is_worker_free', side_effect=free): + resp = client.get('/worker/get', json={ + 'worker_instance_type': 'big', + 'worker_processes': 2, + 'runtime_name': 'python3', + }) + + free_workers = resp.get_json() + assert resp.status_code == 200 + assert [w[0] for w in free_workers] == ['wanted'] + # name, ip, instance id, ssh credentials, instance type, runtime + assert free_workers[0][-1] == 'python3' + + +class TestWorkerRegistration: + + def test_redis_field_flattens_what_a_hash_cannot_hold(self): + assert sa_master._redis_field({'a': 1}) == '{"a": 1}' + assert sa_master._redis_field([1, 2]) == '[1, 2]' + assert sa_master._redis_field(True) == 'True' + assert sa_master._redis_field('text') == 'text' + assert sa_master._redis_field(7) == 7 + + def test_save_worker_keeps_the_backend_section_out_of_redis(self): + worker = _worker_instance() + sa_master.save_worker(worker, _standalone_config(), 'wq:sess-0-M000') + mapping = sa_master.redis_client.hset.call_args.kwargs['mapping'] + assert mapping['name'] == worker.name + assert mapping['status'] == WorkerStatus.STARTING.value + assert mapping['queue_name'] == 'wq:sess-0-M000' + assert mapping['worker_processes'] == 2 + assert 'vm' not in mapping + assert 'do-not-store' not in json.dumps(mapping) + + def test_save_worker_resolves_auto_worker_processes(self): + worker = _worker_instance(config={'worker_processes': 'AUTO'}) + sa_master.save_worker(worker, _standalone_config(), 'wq') + mapping = sa_master.redis_client.hset.call_args.kwargs['mapping'] + assert mapping['worker_processes'] == sa_master.CPU_COUNT + + def test_save_worker_tolerates_an_instance_with_no_ip_yet(self): + worker = _worker_instance(private_ip=None, instance_id=None) + sa_master.save_worker(worker, _standalone_config(), 'wq') + mapping = sa_master.redis_client.hset.call_args.kwargs['mapping'] + assert mapping['private_ip'] == '' + assert mapping['instance_id'] == '' + + def test_worker_vm_data_carries_the_master_and_the_queue(self): + data = sa_master._worker_vm_data(_worker_instance(), 'wq:sess-0-M000') + assert data['master_ip'] == '10.0.0.1' + assert data['work_queue_name'] == 'wq:sess-0-M000' + assert data['lithops_version'] == __version__ + assert data['name'] == 'lithops-worker-1' + + def test_mark_worker_error_records_the_reason(self): + sa_master._mark_worker_error('lithops-worker-1', 'boom') + args, kwargs = sa_master.redis_client.hset.call_args + assert args[0] == 'worker:lithops-worker-1' + assert kwargs['mapping'] == { + 'status': WorkerStatus.ERROR.value, 'err': 'boom' + } + + def test_worker_setup_script_installs_the_host_first(self): + handler = MagicMock() + handler.config = _standalone_config() + with patch.object( + sa_master, 'get_host_setup_script', return_value='HOST;' + ), patch.object( + sa_master, 'get_worker_setup_script', return_value='WORKER;' + ): + script = sa_master._worker_setup_script(handler, {'name': 'w'}) + assert script == 'HOST;WORKER;' + + +class TestWorkerSetup: + + def _handler(self, **cfg): + handler = MagicMock() + handler.config = _standalone_config(**cfg) + return handler + + def test_setup_skips_a_worker_that_is_already_active(self): + handler = self._handler() + handler.backend.get_instance.return_value = _worker_instance() + sa_master.redis_client.hget.return_value = WorkerStatus.ACTIVE.value + sa_master.setup_worker_create_reuse(handler, {'name': 'w'}, 'wq') + sa_master.redis_client.hset.assert_not_called() + + def test_setup_installs_and_leaves_the_worker_installing(self): + handler = self._handler() + worker = _worker_instance() + handler.backend.get_instance.return_value = worker + sa_master.redis_client.hget.return_value = WorkerStatus.STARTING.value + + with patch.object( + sa_master, '_worker_setup_script', return_value='SCRIPT' + ): + sa_master.setup_worker_create_reuse(handler, {'name': 'w'}, 'wq') + + worker.wait_ready.assert_called_once() + worker.validate_capabilities.assert_called_once() + ssh = worker.get_ssh_client.return_value + ssh.upload_local_file.assert_called_once() + ssh.upload_data_to_file.assert_called_once_with( + 'SCRIPT', '/tmp/install_lithops.sh' + ) + # The install runs in the background: the worker reports back itself + assert ssh.run_remote_command.call_args.kwargs['run_async'] is True + worker.del_ssh_client.assert_called_once() + statuses = [ + ( + c.kwargs.get('mapping', {}).get('status') + or (c.args[2] if len(c.args) > 2 else None) + ) + for c in sa_master.redis_client.hset.call_args_list + ] + assert WorkerStatus.INSTALLING.value in statuses + + def test_setup_recreates_a_worker_that_does_not_come_up(self): + handler = self._handler() + worker = _worker_instance(config={ + 'worker_processes': 2, 'worker_create_retries': 2 + }) + worker.wait_ready.side_effect = [TimeoutError('slow'), None] + handler.backend.get_instance.return_value = worker + sa_master.redis_client.hget.return_value = WorkerStatus.STARTING.value + + with patch.object( + sa_master, '_worker_setup_script', return_value='SCRIPT' + ): + sa_master.setup_worker_create_reuse(handler, {'name': 'w'}, 'wq') + + worker.delete.assert_called_once() + worker.create.assert_called_once() + assert worker.wait_ready.call_count == 2 + + def test_setup_gives_up_when_the_worker_never_comes_up(self): + handler = self._handler() + worker = _worker_instance(config={ + 'worker_processes': 2, 'worker_create_retries': 1 + }) + worker.wait_ready.side_effect = TimeoutError('slow') + handler.backend.get_instance.return_value = worker + sa_master.redis_client.hget.return_value = WorkerStatus.STARTING.value + + with pytest.raises(TimeoutError): + sa_master.setup_worker_create_reuse(handler, {'name': 'w'}, 'wq') + + errors = [ + c.kwargs['mapping']['err'] + for c in sa_master.redis_client.hset.call_args_list + if 'mapping' in c.kwargs and 'err' in c.kwargs['mapping'] + and c.kwargs['mapping']['err'] + ] + assert any('Timeout' in e for e in errors) + + def test_setup_records_why_the_installation_failed(self): + handler = self._handler() + worker = _worker_instance() + worker.get_ssh_client.return_value.upload_local_file.side_effect = ( + OSError('no route') + ) + handler.backend.get_instance.return_value = worker + sa_master.redis_client.hget.return_value = WorkerStatus.STARTING.value + + with pytest.raises(OSError): + sa_master.setup_worker_create_reuse(handler, {'name': 'w'}, 'wq') + + errors = [ + c.kwargs['mapping'].get('err') + for c in sa_master.redis_client.hset.call_args_list + if 'mapping' in c.kwargs + ] + assert any(e and 'no route' in e for e in errors) + + def test_consume_setup_runs_the_script_on_this_instance(self, tmp_path): + handler = self._handler(exec_mode=StandaloneMode.CONSUME.value) + instance = _worker_instance() + handler.backend.get_instance.return_value = instance + sa_master.redis_client.hget.return_value = WorkerStatus.STARTING.value + script_path = str(tmp_path / 'install_lithops.sh') + real_open = open + + with ExitStack() as stack: + enter = stack.enter_context + enter(patch.object( + sa_master, '_worker_setup_script', return_value='SCRIPT' + )) + system = enter( + patch.object(sa_master.os, 'system', return_value=0) + ) + chmod = enter(patch.object(sa_master.os, 'chmod')) + remove = enter(patch.object(sa_master.os, 'remove')) + enter(patch( + 'builtins.open', + side_effect=lambda *a, **k: real_open(script_path, 'w'), + )) + sa_master.setup_worker_consume(handler, {'name': 'w'}, 'wq') + + assert instance.private_ip == '10.0.0.1' + system.assert_called_once() + assert system.call_args[0][0].startswith('sudo ') + chmod.assert_called_once() + remove.assert_called_once() + + def test_consume_setup_reports_a_failing_script(self, tmp_path): + handler = self._handler(exec_mode=StandaloneMode.CONSUME.value) + handler.backend.get_instance.return_value = _worker_instance() + sa_master.redis_client.hget.return_value = WorkerStatus.STARTING.value + script_path = str(tmp_path / 'install_lithops.sh') + real_open = open + + with ExitStack() as stack: + enter = stack.enter_context + enter(patch.object( + sa_master, '_worker_setup_script', return_value='SCRIPT' + )) + enter(patch.object(sa_master.os, 'system', return_value=256)) + enter(patch.object(sa_master.os, 'chmod')) + enter(patch.object(sa_master.os, 'remove')) + enter(patch( + 'builtins.open', + side_effect=lambda *a, **k: real_open(script_path, 'w'), + )) + log = enter(patch.object(sa_master, 'logger')) + sa_master.setup_worker_consume(handler, {'name': 'w'}, 'wq') + + assert any( + 'wait status' in str(c) for c in log.error.call_args_list + ) + + +class TestHandleWorkers: + + def test_handle_workers_does_nothing_without_workers(self): + with patch.object(sa_master, 'StandaloneHandler') as handler: + sa_master.handle_workers({'config': {}}, [], 'wq') + handler.assert_not_called() + + def test_handle_workers_sets_each_created_worker_up(self): + payload = {'config': {'standalone': {}}} + with ExitStack() as stack: + enter = stack.enter_context + enter(patch.object( + sa_master, 'extract_standalone_config', + return_value=_standalone_config(), + )) + enter(patch.object(sa_master, 'StandaloneHandler')) + setup = enter( + patch.object(sa_master, 'setup_worker_create_reuse') + ) + sa_master.handle_workers( + payload, [{'name': 'a'}, {'name': 'b'}], 'wq' + ) + assert setup.call_count == 2 + + def test_handle_workers_counts_a_failed_worker_as_one_less(self): + payload = {'config': {'standalone': {}}} + with ExitStack() as stack: + enter = stack.enter_context + enter(patch.object( + sa_master, 'extract_standalone_config', + return_value=_standalone_config(), + )) + enter(patch.object(sa_master, 'StandaloneHandler')) + enter(patch.object( + sa_master, 'setup_worker_create_reuse', + side_effect=[None, OSError('boom')], + )) + log = enter(patch.object(sa_master, 'logger')) + sa_master.handle_workers( + payload, [{'name': 'a'}, {'name': 'b'}], 'wq' + ) + assert log.error.called + assert any('1 of 2' in str(c) for c in log.debug.call_args_list) + + def test_handle_workers_uses_the_master_in_consume_mode(self): + payload = {'config': {'standalone': {}}} + cfg = _standalone_config(exec_mode=StandaloneMode.CONSUME.value) + with ExitStack() as stack: + enter = stack.enter_context + enter(patch.object( + sa_master, 'extract_standalone_config', return_value=cfg + )) + enter(patch.object(sa_master, 'StandaloneHandler')) + consume = enter(patch.object(sa_master, 'setup_worker_consume')) + create = enter( + patch.object(sa_master, 'setup_worker_create_reuse') + ) + sa_master.handle_workers(payload, [{'name': 'a'}], 'wq') + consume.assert_called_once() + create.assert_not_called() + + +def _job_payload(**extra): + payload = { + 'job_key': 'sess-0-M000', + 'executor_id': 'sess-0', + 'job_id': 'M000', + 'host_submit_tstamp': 1700000000.0, + 'func_name': 'add', + 'runtime_name': 'python3', + 'worker_instance_type': 'big', + 'worker_processes': 2, + 'call_ids': ['00000', '00001'], + 'data_byte_ranges': [(0, 9), (10, 19)], + 'worker_instances': [], + 'config': {'standalone': {'exec_mode': StandaloneMode.CREATE.value}}, + } + payload.update(extra) + return payload + + +class TestJobHandling: + + def test_handle_job_registers_it_and_queues_one_task_per_call(self): + sa_master.handle_job(_job_payload(), 'wq:sess-0-M000') + + mapping = sa_master.redis_client.hset.call_args.kwargs['mapping'] + assert mapping['job_key'] == 'sess-0-M000' + assert mapping['status'] == JobStatus.SUBMITTED.value + assert mapping['total_tasks'] == 2 + assert mapping['queue_name'] == 'wq:sess-0-M000' + + pushes = sa_master.redis_client.lpush.call_args_list + assert len(pushes) == 2 + first = json.loads(pushes[0][0][1]) + second = json.loads(pushes[1][0][1]) + # Each task carries only its own call and its own data range + assert first['call_ids'] == ['00000'] + assert first['data_byte_ranges'] == [[0, 9]] + assert second['call_ids'] == ['00001'] + assert second['data_byte_ranges'] == [[10, 19]] + + def test_job_list_builds_a_table_with_the_progress(self, client): + sa_master.redis_client.keys.return_value = ['job:sess-0-M000'] + sa_master.redis_client.hgetall.return_value = { + 'job_key': 'sess-0-M000', + 'status': JobStatus.RUNNING.value, + 'submitted': '1700000000', + 'func_name': 'add', + 'worker_type': 'big', + 'runtime_name': 'python3', + 'exec_mode': StandaloneMode.CREATE.value, + 'total_tasks': '4', + } + sa_master.redis_client.llen.return_value = 3 + table = client.get('/job/list').get_json() + assert table[0][0] == 'Job ID' + assert table[1][0] == 'sess-0-M000' + assert table[1][1] == 'add()' + assert table[1][3] == 'big' + assert table[1][5] == '3/4' + + def test_job_list_calls_the_worker_type_vm_in_consume_mode(self, client): + sa_master.redis_client.keys.return_value = ['job:sess-0-M000'] + sa_master.redis_client.hgetall.return_value = { + 'job_key': 'sess-0-M000', + 'status': JobStatus.RUNNING.value, + 'submitted': '1700000000', + 'func_name': 'add', + 'worker_type': 'ignored', + 'runtime_name': 'python3', + 'exec_mode': StandaloneMode.CONSUME.value, + 'total_tasks': '1', + } + sa_master.redis_client.llen.return_value = 0 + assert client.get('/job/list').get_json()[1][3] == 'VM' + + +class TestJobRun: + + def test_run_rejects_a_body_that_is_not_a_dict(self, client): + resp = client.post('/job/run', json=['nope']) + assert resp.status_code == 404 + + def test_run_rejects_an_invalid_runtime_name(self, client): + resp = client.post( + '/job/run', json=_job_payload(runtime_name='has space') + ) + assert resp.status_code == 404 + + def test_run_accepts_the_job_and_answers_with_an_activation_id( + self, client + ): + with patch.object(sa_master, 'Thread') as thread: + resp = client.post('/job/run', json=_job_payload()) + assert resp.status_code == 202 + assert len(resp.get_json()['activationId']) == 12 + # One thread queues the job, another sets the workers up + assert thread.call_count == 2 + sa_master.budget_keeper.add_job.assert_called_once_with('sess-0-M000') + + def _queue_name_for(self, client, payload): + with patch.object(sa_master, 'Thread') as thread: + client.post('/job/run', json=payload) + return thread.call_args_list[0].kwargs['args'][1] + + def test_create_mode_gives_the_job_its_own_queue(self, client): + payload = _job_payload() + assert self._queue_name_for(client, payload) == 'wq:sess-0-m000' + + def test_consume_mode_queues_by_runtime(self, client): + payload = _job_payload( + runtime_name='lithops/Python:3.12', + config={'standalone': { + 'exec_mode': StandaloneMode.CONSUME.value + }}, + ) + assert self._queue_name_for(client, payload) == ( + 'wq:localhost:lithops-python:3.12' + ) + + def test_reuse_mode_queues_by_worker_shape(self, client): + payload = _job_payload( + config={'standalone': {'exec_mode': StandaloneMode.REUSE.value}} + ) + assert self._queue_name_for(client, payload) == 'wq:big-2-python3' + + def test_run_takes_the_worker_instances_out_of_the_payload(self, client): + payload = _job_payload(worker_instances=[{'name': 'w'}]) + with patch.object(sa_master, 'Thread') as thread: + client.post('/job/run', json=payload) + # The job payload the workers receive carries no instance list + queued_payload = thread.call_args_list[0].kwargs['args'][0] + assert 'worker_instances' not in queued_payload + assert thread.call_args_list[1].kwargs['args'][1] == [{'name': 'w'}] + + +class TestCancelJob: + + def test_cancel_requeues_the_tasks_of_other_jobs(self): + mine = json.dumps({'job_key': 'sess-0-M000'}) + theirs = json.dumps({'job_key': 'sess-0-M001'}) + sa_master.redis_client.hget.return_value = 'wq:sess-0-M000' + sa_master.redis_client.llen.side_effect = [1, 1, 0] + sa_master.redis_client.rpop.side_effect = [mine, theirs] + sa_master.redis_client.keys.return_value = [] + + with patch.object(sa_master.Path, 'touch'): + sa_master.cancel_job_process(['sess-0-M000']) + + pushed = [c[0][1] for c in sa_master.redis_client.lpush.call_args_list] + assert pushed == [theirs] + + def test_cancel_marks_the_job_canceled_and_leaves_a_done_file(self): + sa_master.redis_client.hget.side_effect = [ + 'wq:sess-0-M000', JobStatus.RUNNING.value + ] + sa_master.redis_client.llen.return_value = 0 + sa_master.redis_client.keys.return_value = [] + + with patch.object(sa_master.Path, 'touch') as touch: + sa_master.cancel_job_process(['sess-0-M000']) + + touch.assert_called_once() + assert sa_master.redis_client.hset.call_args[0][2] == ( + JobStatus.CANCELED.value + ) + + def test_cancel_leaves_an_already_done_job_alone(self): + sa_master.redis_client.hget.side_effect = [ + 'wq:sess-0-M000', JobStatus.DONE.value + ] + sa_master.redis_client.llen.return_value = 0 + sa_master.redis_client.keys.return_value = [] + + with patch.object(sa_master.Path, 'touch'): + sa_master.cancel_job_process(['sess-0-M000']) + + sa_master.redis_client.hset.assert_not_called() + + def test_cancel_tells_every_worker_to_stop_the_job(self): + sa_master.redis_client.hget.return_value = 'wq:sess-0-M000' + sa_master.redis_client.llen.return_value = 0 + sa_master.redis_client.keys.return_value = ['worker:w1'] + sa_master.redis_client.hgetall.return_value = { + 'private_ip': '10.0.0.2' + } + + with patch.object(sa_master.Path, 'touch'), \ + patch.object(sa_master.requests, 'post') as post: + sa_master.cancel_job_process(['sess-0-M000']) + + assert '/stop/sess-0-M000' in post.call_args[0][0] + + def test_cancel_reports_a_worker_that_cannot_be_told(self): + sa_master.redis_client.hget.return_value = 'wq:sess-0-M000' + sa_master.redis_client.llen.return_value = 0 + sa_master.redis_client.keys.return_value = ['worker:w1'] + sa_master.redis_client.hgetall.return_value = { + 'private_ip': '10.0.0.2' + } + + with ExitStack() as stack: + enter = stack.enter_context + enter(patch.object(sa_master.Path, 'touch')) + enter(patch.object( + sa_master.requests, 'post', side_effect=OSError('down') + )) + log = enter(patch.object(sa_master, 'logger')) + sa_master.cancel_job_process(['sess-0-M000']) + + assert log.error.called + + +class TestJobMonitor: + + def test_monitor_reports_progress_and_marks_a_job_complete(self): + sa_master.redis_client.keys.return_value = ['job:sess-0-M000'] + sa_master.redis_client.hgetall.return_value = {'total_tasks': '2'} + # One task done on the first pass, both on the second + sa_master.redis_client.llen.side_effect = [1, 2] + + rounds = [] + + def sleep(_seconds): + rounds.append(1) + if len(rounds) > 2: + raise KeyboardInterrupt + + with patch.object(sa_master.time, 'sleep', side_effect=sleep), \ + patch.object(sa_master.Path, 'touch') as touch, \ + patch.object(sa_master, 'logger') as log: + with pytest.raises(KeyboardInterrupt): + sa_master.job_monitor() + + sa_master.budget_keeper.add_job.assert_called_once_with('sess-0-M000') + touch.assert_called_once() + messages = [str(c) for c in log.debug.call_args_list] + assert any('Tasks done: 1/2' in m for m in messages) + assert any('Completed!' in m for m in messages) + + def test_monitor_stops_looking_at_a_job_that_is_done(self): + sa_master.redis_client.keys.return_value = ['job:sess-0-M000'] + sa_master.redis_client.hgetall.return_value = {'total_tasks': '1'} + sa_master.redis_client.llen.side_effect = [1] + + rounds = [] + + def sleep(_seconds): + rounds.append(1) + if len(rounds) > 3: + raise KeyboardInterrupt + + with patch.object(sa_master.time, 'sleep', side_effect=sleep), \ + patch.object(sa_master.Path, 'touch'), \ + patch.object(sa_master, 'logger'): + with pytest.raises(KeyboardInterrupt): + sa_master.job_monitor() + + # llen is only asked while the job still has tasks pending + assert sa_master.redis_client.llen.call_count == 1 + + +class TestMasterMain: + + def test_main_wires_the_keeper_the_monitor_and_the_server(self, tmp_path): + config_file = tmp_path / 'config' + config_file.write_text(json.dumps(_standalone_config())) + data_file = tmp_path / 'master.data' + data_file.write_text(json.dumps({ + 'name': 'lithops-master', 'private_ip': '10.1.2.3' + })) + + keeper = MagicMock() + server = MagicMock() + with ExitStack() as stack: + enter = stack.enter_context + enter(patch.object( + sa_master, 'SA_CONFIG_FILE', str(config_file) + )) + enter(patch.object( + sa_master, 'SA_MASTER_DATA_FILE', str(data_file) + )) + enter(patch.object(sa_master, '_configure_logging')) + keeper_cls = enter(patch.object( + sa_master, 'BudgetKeeper', return_value=keeper + )) + redis_cls = enter(patch.object(sa_master.redis, 'Redis')) + thread = enter(patch.object(sa_master, 'Thread')) + server_cls = enter(patch.object( + sa_master, 'WSGIServer', return_value=server + )) + sa_master.main() + + assert sa_master.master_ip == '10.1.2.3' + keeper_cls.assert_called_once() + keeper.start.assert_called_once() + redis_cls.assert_called_once_with(decode_responses=True) + # The job monitor runs in the background, the server in the foreground + assert thread.call_args.kwargs['target'] is sa_master.job_monitor + assert thread.call_args.kwargs['daemon'] is True + assert server_cls.call_args[0][0] == ( + '0.0.0.0', sa_master.SA_MASTER_SERVICE_PORT + ) + server.serve_forever.assert_called_once() diff --git a/lithops/tests/test_storage_layer.py b/lithops/tests/test_storage_layer.py new file mode 100644 index 000000000..99551fb98 --- /dev/null +++ b/lithops/tests/test_storage_layer.py @@ -0,0 +1,430 @@ +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import json +import pickle +import threading +from unittest.mock import MagicMock, patch + +import pytest + +from lithops.constants import JOBS_PREFIX, RUNTIMES_PREFIX, TEMP_PREFIX +from lithops.storage.cloud_proxy import ( + CloudFileProxy, + CloudStorage, + DelayedBytesBuffer, + DelayedStringBuffer, + _path, + cloud_open, + remove_lithops_keys, +) +from lithops.storage.storage import ( + RUNTIME_META_CACHE, + InternalStorage, + Storage, +) +from lithops.storage.utils import ( + CloudObject, + CloudObjectLocal, + CloudObjectUrl, + StorageConfigMismatchError, + StorageNoSuchKeyError, + check_storage_path, + clean_bucket, + create_data_key, + create_func_key, + create_init_key, + create_job_key, + create_output_key, + create_status_key, + get_storage_path, +) + + +def _bare_storage(backend='localhost'): + storage = Storage.__new__(Storage) + storage.backend = backend + storage.bucket = 'storage' + storage.config = {'backend': backend, backend: {'storage_bucket': 'storage'}} + storage.storage_handler = MagicMock() + return storage + + +def _bare_internal(bucket='storage'): + internal = InternalStorage.__new__(InternalStorage) + internal.storage = MagicMock() + internal.backend = 'localhost' + internal.bucket = bucket + return internal + + +class FakeCloudStorage: + def __init__(self, keys=None, data=None): + self.keys = list(keys or []) + self.data = dict(data or {}) + self.deleted = [] + self.puts = [] + + def list_bucket_keys(self, prefix=None): + if prefix is None: + return list(self.keys) + return [key for key in self.keys if key.startswith(prefix)] + + def delete_data(self, key): + self.deleted.append(key) + + def get_data(self, key): + return self.data[key] + + def put_data(self, key, data): + self.puts.append((key, data)) + self.data[key] = data + + +class TestStorageUtils: + + def test_exception_messages(self): + err = StorageNoSuchKeyError('bucket', 'key') + assert str(err) == 'No such key /bucket/key found in storage.' + mismatch = StorageConfigMismatchError(['a', 'b'], ['c', 'd']) + assert 'stored at' in str(mismatch) + assert "['c', 'd']" in str(mismatch) + assert "['a', 'b']" in str(mismatch) + + def test_cloudobject_str_forms(self): + assert str(CloudObject('localhost', 'b', 'k')) == ( + '' + ) + assert str(CloudObjectUrl('https://x')) == '' + local = CloudObjectLocal('/tmp/dir/file.txt') + assert local.bucket == '/tmp/dir' + assert local.key == 'file.txt' + assert str(local) == '' + + def test_job_key_builders(self): + assert create_job_key('sess-0', 'M000') == 'sess-0-M000' + assert create_func_key('sess-0', 'abc') == ( + f'{JOBS_PREFIX}/sess-0/abc.func.pickle' + ) + assert create_data_key('sess-0', 'M000') == ( + f'{JOBS_PREFIX}/sess-0-M000/aggdata.pickle' + ) + assert create_output_key('sess-0', 'M000', '00000') == ( + f'{JOBS_PREFIX}/sess-0-M000/00000/output.pickle' + ) + assert create_status_key('sess-0', 'M000', '00000') == ( + f'{JOBS_PREFIX}/sess-0-M000/00000/status.json' + ) + assert create_init_key('sess-0', 'M000', '00000', 'act') == ( + f'{JOBS_PREFIX}/sess-0-M000/00000/act.init' + ) + for key in ( + create_func_key('sess-0', 'abc'), + create_data_key('sess-0', 'M000'), + create_output_key('sess-0', 'M000', '00000'), + ): + assert '\\' not in key + assert key.startswith(f'{JOBS_PREFIX}/') + + def test_storage_path_is_a_list_and_check_raises(self): + cfg = {'backend': 'localhost', 'localhost': {'storage_bucket': 'b'}} + path = get_storage_path(cfg) + assert path == ['localhost', 'b'] + assert type(path) is list + check_storage_path(cfg, ['localhost', 'b']) + with pytest.raises(StorageConfigMismatchError): + check_storage_path(cfg, ['s3', 'other']) + + def test_clean_bucket_loops_until_empty(self): + storage = MagicMock() + storage.list_keys.side_effect = [['a', 'b'], ['c'], []] + sleeps = [] + test_thread = threading.current_thread() + + def sleep(seconds): + if threading.current_thread() is test_thread: + sleeps.append(seconds) + + with patch('lithops.storage.utils.time.sleep', side_effect=sleep): + clean_bucket(storage, 'bucket', 'pref', sleep=2) + assert storage.delete_objects.call_count == 2 + assert sleeps == [2, 2] + + +class TestStorageCloudObjects: + + def test_put_cloudobject_uses_temp_prefix_and_hex_id(self, monkeypatch): + monkeypatch.delenv('__LITHOPS_SESSION_ID', raising=False) + storage = _bare_storage() + cloudobject = storage.put_cloudobject(b'data') + key = storage.storage_handler.put_object.call_args[0][1] + assert key.startswith(TEMP_PREFIX + '/') + assert 'cloudobject_' in key + assert cloudobject.backend == 'localhost' + assert cloudobject.bucket == 'storage' + + def test_put_cloudobject_prefixes_session_id(self, monkeypatch): + monkeypatch.setenv('__LITHOPS_SESSION_ID', 'sess') + storage = _bare_storage() + cloudobject = storage.put_cloudobject(b'data') + key = storage.storage_handler.put_object.call_args[0][1] + assert '/sess/cloudobject_' in key + assert cloudobject.key == key + + def test_get_and_delete_cloudobject_reject_other_backend(self): + storage = _bare_storage() + other = CloudObject('s3', 'b', 'k') + with pytest.raises(Exception, match='Invalid Storage backend'): + storage.get_cloudobject(other) + with pytest.raises(Exception, match='Invalid Storage backend'): + storage.delete_cloudobject(other) + own = CloudObject('localhost', 'b', 'k') + storage.get_cloudobject(own, stream=True) + storage.storage_handler.get_object.assert_called_once_with( + 'b', 'k', stream=True + ) + + def test_delete_cloudobjects_groups_by_bucket(self): + storage = _bare_storage() + storage.delete_cloudobjects([ + CloudObject('localhost', 'b1', 'k1'), + CloudObject('localhost', 'b2', 'k2'), + CloudObject('localhost', 'b1', 'k3'), + ]) + calls = storage.storage_handler.delete_objects.call_args_list + deleted = {call[0][0]: set(call[0][1]) for call in calls} + assert deleted == {'b1': {'k1', 'k3'}, 'b2': {'k2'}} + + def test_delete_cloudobjects_rejects_other_backend(self): + storage = _bare_storage() + with pytest.raises(Exception, match='Invalid Storage backend'): + storage.delete_cloudobjects([CloudObject('s3', 'b', 'k')]) + + def test_create_bucket_skips_when_backend_has_no_method(self): + storage = _bare_storage() + del storage.storage_handler.create_bucket + assert storage.create_bucket('b') is None + + def test_get_object_keeps_mutable_extra_get_args_default(self): + assert Storage.get_object.__defaults__[-1] == {} + + +class TestInternalStorage: + + def test_missing_bucket_raises(self): + storage = MagicMock() + storage.backend = 'localhost' + storage.bucket = None + with patch('lithops.storage.storage.Storage', return_value=storage): + with pytest.raises(Exception, match='storage_bucket'): + InternalStorage({'backend': 'localhost'}) + + def test_get_job_status_parses_init_and_status_keys(self): + internal = _bare_internal() + keys = [ + create_init_key('sess-0', 'M000', '00000', 'act1'), + create_status_key('sess-0', 'M000', '00001'), + f'{JOBS_PREFIX}/ignored.txt', + ] + internal.storage.list_keys.return_value = keys + running, done = internal.get_job_status('sess-0') + assert (('sess-0', 'M000', '00000'), 'act1') in running + assert ('sess-0', 'M000', '00001') in done + internal.storage.list_keys.assert_called_once_with( + 'storage', f'{JOBS_PREFIX}/sess-0' + ) + + def test_get_job_status_lists_per_job_when_job_ids_given(self): + internal = _bare_internal() + keys = [ + create_init_key('sess-0', 'M000', '00000', 'act1'), + create_status_key('sess-0', 'M000', '00001'), + ] + internal.storage.list_keys.return_value = keys + running, done = internal.get_job_status('sess-0', job_ids=['M000']) + assert (('sess-0', 'M000', '00000'), 'act1') in running + assert ('sess-0', 'M000', '00001') in done + internal.storage.list_keys.assert_called_once_with( + 'storage', f'{JOBS_PREFIX}/{create_job_key("sess-0", "M000")}' + ) + + def test_get_call_status_and_output_missing_are_none(self): + internal = _bare_internal() + internal.storage.get_object.side_effect = StorageNoSuchKeyError('b', 'k') + assert internal.get_call_status('e', 'j', 'c') is None + assert internal.get_call_output('e', 'j', 'c') is None + + def test_get_call_status_decodes_ascii_json(self): + internal = _bare_internal() + internal.storage.get_object.return_value = b'{"ok": true}' + assert internal.get_call_status('e', 'M000', '00000') == {'ok': True} + + def test_runtime_meta_memory_cache(self): + RUNTIME_META_CACHE.clear() + internal = _bare_internal() + cache_key = f'{RUNTIMES_PREFIX}/rk.meta.json' + RUNTIME_META_CACHE[cache_key] = {'cached': True} + assert internal.get_runtime_meta('rk') == {'cached': True} + internal.storage.get_object.assert_not_called() + RUNTIME_META_CACHE.clear() + + def test_runtime_meta_disk_then_storage_then_missing(self, tmp_path, monkeypatch): + RUNTIME_META_CACHE.clear() + monkeypatch.setattr('lithops.storage.storage.CACHE_DIR', str(tmp_path)) + monkeypatch.setattr( + 'lithops.storage.storage.is_lithops_worker', lambda: False + ) + internal = _bare_internal() + meta_dir = tmp_path / RUNTIMES_PREFIX + meta_dir.mkdir() + (meta_dir / 'rk.meta.json').write_text(json.dumps({'from': 'disk'})) + assert internal.get_runtime_meta('rk') == {'from': 'disk'} + + RUNTIME_META_CACHE.clear() + (meta_dir / 'rk.meta.json').unlink() + internal.storage.get_object.return_value = b'{"from": "storage"}' + assert internal.get_runtime_meta('rk') == {'from': 'storage'} + assert (meta_dir / 'rk.meta.json').exists() + + RUNTIME_META_CACHE.clear() + (meta_dir / 'rk.meta.json').unlink() + internal.storage.get_object.side_effect = StorageNoSuchKeyError('b', 'k') + assert internal.get_runtime_meta('rk') is None + RUNTIME_META_CACHE.clear() + + def test_put_and_delete_runtime_meta(self, tmp_path, monkeypatch): + monkeypatch.setattr('lithops.storage.storage.CACHE_DIR', str(tmp_path)) + monkeypatch.setattr( + 'lithops.storage.storage.is_lithops_worker', lambda: False + ) + internal = _bare_internal() + internal.put_runtime_meta('rk', {'a': 1}) + obj_key = f'{RUNTIMES_PREFIX}/rk.meta.json' + internal.storage.put_object.assert_called_once() + assert internal.storage.put_object.call_args[0][1] == obj_key + local = tmp_path / RUNTIMES_PREFIX / 'rk.meta.json' + assert json.loads(local.read_text()) == {'a': 1} + + internal.delete_runtime_meta('rk') + assert not local.exists() + internal.storage.delete_object.assert_called_once_with('storage', obj_key) + + +class TestCloudProxy: + + def test_remove_lithops_keys(self): + keys = [ + f'{JOBS_PREFIX}/a', + 'user/file', + f'{TEMP_PREFIX}/x', + f'{RUNTIMES_PREFIX}/r', + 'other', + ] + assert remove_lithops_keys(keys) == ['user/file', 'other'] + + def test_listdir_and_path_helpers(self): + fake = FakeCloudStorage(keys=[ + 'dir/a.txt', + 'dir/sub/b.txt', + f'{JOBS_PREFIX}/hidden', + 'dir/c.txt', + ]) + proxy = CloudFileProxy(fake) + listed = proxy.listdir('dir', suffix_dirs=True) + assert 'a.txt' in listed + assert 'c.txt' in listed + assert 'sub/' in listed + assert not any(name.startswith(JOBS_PREFIX) for name in listed) + + assert proxy.path.isfile('dir/a.txt') is True + assert proxy.path.isdir('dir') is True + assert proxy.path.exists('dir/a.txt') is True + assert proxy.path.exists('missing') is False + + def test_exists_keeps_leading_slash_on_list_prefix(self): + fake = FakeCloudStorage() + fake.list_bucket_keys = MagicMock(return_value=[]) + _path(fake).exists('/foo') + fake.list_bucket_keys.assert_called_once_with(prefix='/foo') + + def test_listdir_of_the_root_matches_the_slash_form(self): + # The empty path is the default argument, and it used to ask for the + # prefix '/', which no key can start with + fake = FakeCloudStorage(keys=['top.txt', 'dir/inner.txt']) + assert sorted(CloudFileProxy(fake).listdir('')) == ['dir', 'top.txt'] + assert sorted(CloudFileProxy(fake).listdir('/')) == ['dir', 'top.txt'] + + def test_walk_yields_nothing_when_missing(self): + # os.walk yields nothing for a path that is not there, and raising + # StopIteration inside a generator only became a RuntimeError + proxy = CloudFileProxy(FakeCloudStorage(keys=[])) + assert list(proxy.walk('missing')) == [] + + def test_open_rejects_unsupported_mode(self): + fake = FakeCloudStorage(data={'f.txt': b'hello'}) + with pytest.raises(ValueError, match='Unsupported mode'): + cloud_open('f.txt', mode='a', cloud_storage=fake) + + def test_open_read_and_write_buffers(self): + fake = FakeCloudStorage(data={'f.txt': b'hello'}) + text = cloud_open('f.txt', mode='r', cloud_storage=fake) + assert text.read() == 'hello' + binary = cloud_open('f.txt', mode='rb', cloud_storage=fake) + assert binary.read() == b'hello' + + buf = cloud_open('out.txt', mode='w', cloud_storage=fake) + buf.write('world') + buf.close() + assert fake.puts[-1] == ('out.txt', 'world') + + bbuf = cloud_open('out.bin', mode='wb', cloud_storage=fake) + bbuf.write(b'xyz') + bbuf.close() + assert fake.puts[-1] == ('out.bin', b'xyz') + + def test_delayed_buffers_run_action_on_close(self): + seen = [] + DelayedBytesBuffer(seen.append, b'ab').close() + DelayedStringBuffer(seen.append, 'cd').close() + assert seen == [b'ab', 'cd'] + + def test_remove_delegates_and_mkdir_is_noop(self): + fake = FakeCloudStorage() + proxy = CloudFileProxy(fake) + proxy.remove('x') + assert fake.deleted == ['x'] + proxy.mkdir('whatever') + proxy.makedirs('whatever') + + def test_cloud_storage_config_branches(self): + with patch.object(Storage, '__init__', return_value=None): + raw = {'backend': 'localhost', 'localhost': {'storage_bucket': 'b'}} + assert CloudStorage(raw)._config is raw + extracted = {'backend': 'localhost', 'localhost': {'storage_bucket': 'x'}} + with patch( + 'lithops.storage.cloud_proxy.extract_storage_config', + return_value=extracted, + ): + lithops_cfg = {'lithops': {}, 'localhost': {}} + assert CloudStorage(lithops_cfg)._config is extracted + + def test_cloud_storage_pickle_roundtrip_uses_config(self): + cfg = {'backend': 'localhost', 'localhost': {'storage_bucket': 'storage'}} + with patch.object(Storage, '__init__', return_value=None): + original = CloudStorage(cfg) + dumped = pickle.dumps(original) + with patch.object(Storage, '__init__', return_value=None) as init: + loaded = pickle.loads(dumped) + assert loaded._config == cfg + init.assert_called() diff --git a/lithops/tests/test_util.py b/lithops/tests/test_util.py new file mode 100644 index 000000000..476a6d01d --- /dev/null +++ b/lithops/tests/test_util.py @@ -0,0 +1,418 @@ +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import os +import pickle +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from lithops.util.ibm_token_manager import ( + COSTokenManager, + EXPIRY_MINUTES, + IAMTokenManager, + IBMTokenManager, +) +from lithops.util.metrics import PrometheusExporter +from lithops.util.ssh_client import SSHClient, ssh_boot_status_message + + +class TestSshBootStatusMessage: + + def test_timeout(self): + assert 'waiting for network' in ssh_boot_status_message( + TimeoutError('timed out') + ) + + def test_connection_refused(self): + assert 'starting SSH' in ssh_boot_status_message( + OSError('Connection refused') + ) + + def test_banner(self): + assert 'Configuring SSH' in ssh_boot_status_message( + Exception('Error reading SSH protocol banner') + ) + + def test_fallback_to_str(self): + assert ssh_boot_status_message(Exception('weird')) == 'weird' + + +class TestSSHClient: + + def test_expands_key_filename(self, tmp_path): + key = tmp_path / 'id_rsa' + key.write_text('k') + creds = {'username': 'u', 'key_filename': str(key)} + SSHClient('1.2.3.4', creds) + assert creds['key_filename'] == str(key) + + def test_missing_key_falls_back_to_default(self, tmp_path): + creds = {'username': 'u', 'key_filename': str(tmp_path / 'missing')} + SSHClient('1.2.3.4', creds) + assert creds['key_filename'] == os.path.expanduser('~/.ssh/id_rsa') + + def test_invalid_ip_raises(self): + client = SSHClient('0.0.0.0', {}) + with pytest.raises(Exception, match='Invalid IP Address'): + client.run_remote_command('true') + + def test_create_client_passes_key_filename(self, tmp_path): + key = tmp_path / 'id_rsa' + key.write_text('k') + creds = { + 'username': 'ubuntu', + 'password': None, + 'key_filename': str(key), + } + client = SSHClient('10.0.0.1', creds) + ssh = MagicMock() + with patch('lithops.util.ssh_client.paramiko.SSHClient', return_value=ssh): + client.create_client(timeout=5) + kwargs = ssh.connect.call_args.kwargs + assert kwargs['hostname'] == '10.0.0.1' + assert kwargs['username'] == 'ubuntu' + assert kwargs['key_filename'] == str(key) + assert 'pkey' not in kwargs + + def test_run_remote_retries_on_exec_failure(self): + client = SSHClient('10.0.0.1', {'username': 'u'}) + ssh = MagicMock() + stdout = MagicMock() + stdout.read.return_value = b'ok\n' + stderr = MagicMock() + stderr.read.return_value = b'' + ssh.exec_command.side_effect = [ + Exception('timeout'), + (MagicMock(), stdout, stderr), + ] + with patch.object(client, 'create_client', return_value=ssh) as created: + client.ssh_client = ssh + out, err = client.run_remote_command('echo ok') + assert created.call_count == 1 + assert out == 'ok' + assert err == '' + + def test_download_creates_parent_dir(self, tmp_path): + dest = tmp_path / 'nested' / 'file.txt' + client = SSHClient('10.0.0.1', {}) + ftp = MagicMock() + ssh = MagicMock() + ssh.open_sftp.return_value = ftp + client.ssh_client = ssh + client.download_remote_file('/remote', str(dest)) + assert dest.parent.is_dir() + ftp.get.assert_called_once_with('/remote', str(dest)) + ftp.close.assert_called_once() + + def test_sftp_closes_on_error(self): + client = SSHClient('10.0.0.1', {}) + ftp = MagicMock() + ftp.put.side_effect = OSError('fail') + ssh = MagicMock() + ssh.open_sftp.return_value = ftp + client.ssh_client = ssh + with pytest.raises(OSError): + client.upload_local_file('/local', '/remote') + ftp.close.assert_called_once() + + def test_upload_multiple_and_data(self): + client = SSHClient('10.0.0.1', {}) + ftp = MagicMock() + remote = MagicMock() + ftp.open.return_value.__enter__.return_value = remote + ssh = MagicMock() + ssh.open_sftp.return_value = ftp + client.ssh_client = ssh + client.upload_multiple_local_files([('a', 'b'), ('c', 'd')]) + assert ftp.put.call_count == 2 + client.upload_data_to_file('hello', '/dst') + remote.write.assert_called_once_with('hello') + + +class TestPrometheusExporter: + + def test_missing_session_id_does_not_raise(self, monkeypatch): + monkeypatch.delenv('__LITHOPS_SESSION_ID', raising=False) + exporter = PrometheusExporter(False, None) + assert exporter.instance == 'lithops' + exporter.send_metric('n', 1, type='gauge', labels=[]) + + def test_instance_from_session_id(self, monkeypatch): + monkeypatch.setenv('__LITHOPS_SESSION_ID', 'ek-j0-00000') + exporter = PrometheusExporter(True, {'apigateway': 'http://prom'}) + assert exporter.instance == 'ek' + + def test_send_metric_posts_when_enabled(self, monkeypatch): + monkeypatch.setenv('__LITHOPS_SESSION_ID', 'sid-1') + exporter = PrometheusExporter(True, {'apigateway': 'http://prom'}) + with patch('lithops.util.metrics.requests.post') as post: + exporter.send_metric( + 'function_start', 1.5, type='gauge', + labels=[('job_id', 'j'), ('call_id', 'c')], + ) + post.assert_called_once() + url = post.call_args[0][0] + assert url.startswith('http://prom/metrics/') + assert 'job/lithops' in url + assert 'function_start' in post.call_args.kwargs['data'] + + def test_send_metric_swallows_post_errors(self, monkeypatch): + monkeypatch.setenv('__LITHOPS_SESSION_ID', 'sid-1') + exporter = PrometheusExporter(True, {'apigateway': 'http://prom'}) + with patch( + 'lithops.util.metrics.requests.post', side_effect=OSError('down') + ): + exporter.send_metric('n', 1, type='gauge', labels=[]) + + +class _StubTokenManager(IBMTokenManager): + TOKEN_FILE = None + TYPE = 'TEST' + + def _generate_new_token(self): + self.token = 'new-token' + self.expiry_time = int( + (datetime.now(timezone.utc) + timedelta(hours=1)).timestamp() + ) + + +class TestIBMTokenManager: + + def test_token_file_constant_is_spelled_correctly(self): + assert hasattr(COSTokenManager, 'TOKEN_FILE') + assert hasattr(IAMTokenManager, 'TOKEN_FILE') + assert 'ibm_cos' in COSTokenManager.TOKEN_FILE + assert 'ibm_iam' in IAMTokenManager.TOKEN_FILE + assert not hasattr(COSTokenManager, 'TOEKN_FILE') + + def test_missing_expiry_is_expired(self): + mgr = _StubTokenManager('key') + assert mgr._get_token_minutes_left() == 0 + assert mgr._is_token_expired() + + def test_reuses_unexpired_token(self): + expiry = int( + (datetime.now(timezone.utc) + timedelta(hours=2)).timestamp() + ) + mgr = _StubTokenManager('key', token='cached', token_expiry_time=expiry) + assert mgr._get_token_minutes_left() >= EXPIRY_MINUTES + token, exp = mgr.get_token() + assert token == 'cached' + assert exp == expiry + + def test_refresh_dumps_and_returns_new_token(self, tmp_path, monkeypatch): + path = tmp_path / 'token' + monkeypatch.setattr(_StubTokenManager, 'TOKEN_FILE', str(path)) + mgr = _StubTokenManager('key') + with patch( + 'lithops.util.ibm_token_manager.dump_yaml_config' + ) as dump: + token, expiry = mgr.refresh_token() + assert token == 'new-token' + assert expiry + dump.assert_called_once() + assert dump.call_args[0][0] == str(path) + + def test_loads_cache_file(self, tmp_path, monkeypatch): + path = tmp_path / 'token' + monkeypatch.setattr(_StubTokenManager, 'TOKEN_FILE', str(path)) + path.write_text('x') + expiry = int( + (datetime.now(timezone.utc) + timedelta(hours=2)).timestamp() + ) + with patch( + 'lithops.util.ibm_token_manager.load_yaml_config', + return_value={'token': 'from-disk', 'expiry_time': expiry}, + ): + with patch( + 'lithops.util.ibm_token_manager.os.path.exists', + return_value=True, + ): + mgr = _StubTokenManager('key') + assert mgr.token == 'from-disk' + assert mgr.expiry_time == expiry + + +class TestJoblibBackend: + + def test_consider_sharing_and_handle_call(self): + pytest.importorskip('joblib') + numpy = pytest.importorskip('numpy') + from lithops.util.joblib.lithops_backend import ( + consider_sharing, + handle_call_process, + ) + assert consider_sharing([1, 2]) + assert consider_sharing(numpy.array([1])) + assert not consider_sharing({'a': 1}) + assert handle_call_process(lambda x: x + 1, (2,), {}) == 3 + assert handle_call_process(lambda **kw: kw['v'], (), {'v': 9}) == 9 + + def test_find_shared_objects_proxies_repeated_lists(self): + pytest.importorskip('joblib') + pytest.importorskip('numpy') + from lithops.util.joblib.lithops_backend import find_shared_objects + + shared = [1, 2, 3] + calls = [ + (None, (shared,), {}), + (None, (shared,), {'k': shared}), + ] + storage = MagicMock() + storage.put_cloudobject.return_value = 'cloud-obj' + with patch( + 'lithops.util.joblib.lithops_backend.Storage', + return_value=storage, + ): + out = find_shared_objects(calls) + assert out[0][1][0] == 'cloud-obj' + assert out[1][2]['k'] == 'cloud-obj' + assert 0 in out[0][3] + storage.put_cloudobject.assert_called_once() + + def test_submit_is_the_hook_joblib_calls(self): + # joblib renamed the hook from apply_async to submit, and the + # multiprocessing backend Lithops extends carries its own submit. If + # the override goes away joblib silently stops using this backend + pytest.importorskip('joblib') + pytest.importorskip('diskcache') + from joblib._parallel_backends import PoolManagerMixin + from lithops.util.joblib.lithops_backend import LithopsBackend + + assert 'submit' in LithopsBackend.__dict__ + assert LithopsBackend.submit is not PoolManagerMixin.submit + # Older joblib calls the old name + assert LithopsBackend.apply_async is LithopsBackend.submit + + def test_submit_optimizes_the_batch_before_queueing_it(self): + pytest.importorskip('joblib') + pytest.importorskip('diskcache') + from lithops.util.joblib.lithops_backend import LithopsBackend + + backend = LithopsBackend.__new__(LithopsBackend) + backend.prefer = None + batch = SimpleNamespace(items=[(print, (1,), {})]) + pool = MagicMock() + optimizer = patch( + 'lithops.util.joblib.lithops_backend.find_shared_objects', + return_value=['optimized'], + ) + with patch.object(LithopsBackend, '_get_pool', return_value=pool), \ + optimizer as optimize: + backend.submit(batch, callback='cb') + + optimize.assert_called_once_with(batch.items) + pool.starmap_async.assert_called_once() + assert pool.starmap_async.call_args[0][1] == ['optimized'] + + def test_submit_runs_the_batch_in_one_call_when_threads_preferred(self): + pytest.importorskip('joblib') + pytest.importorskip('diskcache') + from lithops.util.joblib.lithops_backend import LithopsBackend + + backend = LithopsBackend.__new__(LithopsBackend) + backend.prefer = 'threads' + batch = SimpleNamespace(items=[(print, (1,), {})]) + pool = MagicMock() + optimizer = patch( + 'lithops.util.joblib.lithops_backend.find_shared_objects', + return_value=['optimized'], + ) + with patch.object(LithopsBackend, '_get_pool', return_value=pool), \ + optimizer: + backend.submit(batch) + pool.apply_async.assert_called_once() + pool.starmap_async.assert_not_called() + + def test_proxied_args_are_read_from_the_cache_in_one_call(self): + # The tasks of a runtime share the cache directory, and a value too + # big to sit inline is a file of its own, so a key can be present + # while its file is not readable yet. Asking and then reading raised + # KeyError out of the worker + pytest.importorskip('joblib') + pytest.importorskip('diskcache') + from lithops.util.joblib.lithops_backend import replace_with_values + + class RowWithoutItsFileYet: + def __contains__(self, key): + return True + + def __getitem__(self, key): + raise KeyError(key) + + def get(self, key, default=None): + return default + + def __setitem__(self, key, value): + pass + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + storage = MagicMock() + storage.get_cloudobject.return_value = pickle.dumps([1, 2, 3]) + with patch( + 'lithops.util.joblib.lithops_backend.diskcache.Cache', + return_value=RowWithoutItsFileYet(), + ), patch( + 'lithops.util.joblib.lithops_backend.Storage', + return_value=storage, + ): + args, kwargs = replace_with_values(('cloud-obj',), {}, [0]) + + assert args == [[1, 2, 3]] + storage.get_cloudobject.assert_called_once_with('cloud-obj') + + def test_proxied_args_come_from_the_cache_when_it_has_them(self): + pytest.importorskip('joblib') + pytest.importorskip('diskcache') + from lithops.util.joblib.lithops_backend import replace_with_values + + cache = MagicMock() + cache.__enter__ = lambda self: self + cache.__exit__ = lambda self, *exc: False + cache.get.return_value = [4, 5] + with patch( + 'lithops.util.joblib.lithops_backend.diskcache.Cache', + return_value=cache, + ), patch( + 'lithops.util.joblib.lithops_backend.Storage' + ) as storage_cls: + args, kwargs = replace_with_values((), {'k': 'cloud-obj'}, ['k']) + + assert kwargs == {'k': [4, 5]} + storage_cls.assert_not_called() + + def test_find_shared_objects_skips_unique_args(self): + pytest.importorskip('joblib') + pytest.importorskip('numpy') + from lithops.util.joblib.lithops_backend import find_shared_objects + + calls = [ + (None, ([1],), {}), + (None, ([2],), {}), + ] + with patch( + 'lithops.util.joblib.lithops_backend.Storage' + ) as storage_cls: + out = find_shared_objects(calls) + storage_cls.assert_not_called() + assert out[0][1][0] == [1] + assert out[1][1][0] == [2] diff --git a/lithops/tests/test_utils.py b/lithops/tests/test_utils.py new file mode 100644 index 000000000..86571790d --- /dev/null +++ b/lithops/tests/test_utils.py @@ -0,0 +1,658 @@ +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import io +import logging +import pickle +import zipfile +from collections import namedtuple +from unittest.mock import MagicMock, patch + +import pytest + +from lithops import constants +from lithops.utils import ( + CountDownLatch, + CURRENT_PY_VERSION, + FuturesList, + WrappedStreamingBody, + WrappedStreamingBodyPartition, + _as_future_list, + _future_id, + agg_data, + b64str_to_bytes, + b64str_to_dict, + bytes_to_b64str, + convert_bools_to_string, + create_executor_id, + create_futures_list, + dict_to_b64str, + docker_login, + find_free_port, + format_data, + get_default_backend, + get_default_container_name, + get_docker_path, + get_executor_id, + get_mode, + is_linux_system, + is_lithops_worker, + is_notebook, + is_object_processing_function, + is_podman, + is_unix_system, + iterchunks, + MONITORING_QUEUES_ENV, + monitoring_queue_name, + monitoring_queues, + log_prefix, + run_command, + sdb_to_dict, + setup_lithops_logger, + ShutdownSafeStreamHandler, + sizeof_fmt, + split_object_url, + split_path, + timeout_handler, + verify_args, + verify_runtime_name, + version_str, +) + + +class TestGetModeAndBackend: + + def test_get_default_backend_known_modes(self): + assert get_default_backend(constants.LOCALHOST) == constants.LOCALHOST + assert get_default_backend(constants.SERVERLESS) == constants.SERVERLESS_BACKEND_DEFAULT + assert get_default_backend(constants.STANDALONE) == constants.STANDALONE_BACKEND_DEFAULT + + def test_get_default_backend_falsy_mode_returns_none(self): + assert get_default_backend(None) is None + assert get_default_backend('') is None + assert get_default_backend(0) is None + + def test_get_default_backend_unknown_mode_keeps_historical_typo(self): + with pytest.raises(Exception, match='Unknown execution mode: mystery'): + get_default_backend('mystery') + + def test_get_mode_none_uses_default(self): + assert get_mode(None) == constants.MODE_DEFAULT + + def test_get_mode_known_backends(self): + assert get_mode(constants.LOCALHOST) == constants.LOCALHOST + assert get_mode(constants.SERVERLESS_BACKEND_DEFAULT) == constants.SERVERLESS + assert get_mode(constants.STANDALONE_BACKEND_DEFAULT) == constants.STANDALONE + + def test_get_mode_falsy_unknown_returns_none(self): + assert get_mode('') is None + + def test_get_mode_unknown_backend_raises(self): + with pytest.raises(Exception, match='Unknown compute backend: mystery'): + get_mode('mystery') + + +class TestFormatData: + + def test_wraps_scalar_and_converts_range_and_set(self): + assert format_data(7, None) == [7] + assert format_data(range(3), None) == [0, 1, 2] + assert set(format_data({1, 2}, None)) == {1, 2} + + def test_keeps_list_and_futures_list_identity(self): + data = [1, 2] + futures = FuturesList([object(), object()]) + assert format_data(data, None) is data + assert format_data(futures, None) is futures + + def test_tuple_extra_args_concatenated(self): + assert format_data([(1,), (2,)], (10,)) == [(1, 10), (2, 10)] + + def test_tuple_extra_args_must_be_tuple(self): + with pytest.raises(Exception, match='extra_args must contain args in a tuple'): + format_data([(1,)], [10]) + + def test_dict_extra_args_merged_in_place(self): + first = {'a': 1} + result = format_data([first], {'b': 2}) + assert result == [{'a': 1, 'b': 2}] + assert first is result[0] + + def test_dict_extra_args_must_be_dict(self): + with pytest.raises(Exception, match='extra_args must contain kwargs in a dictionary'): + format_data([{'a': 1}], ('b',)) + + def test_scalar_plus_extra_args_becomes_tuple(self): + assert format_data([1, 2], (9, 8)) == [(1, 9, 8), (2, 9, 8)] + + def test_namedtuple_plus_extra_args_is_not_concatenated(self): + Point = namedtuple('Point', 'x y') + pt = Point(1, 2) + # Historical: `type(namedtuple) is tuple` is false, so extra_args are + # wrapped with the namedtuple instead of concatenated. + assert format_data([pt], (9,)) == [(pt, 9)] + + +class TestVerifyArgs: + + def test_futures_list_becomes_future_kwargs(self): + futures = FuturesList(['f1', 'f2']) + assert verify_args(lambda x: x, futures, None) == [ + {'future': 'f1'}, + {'future': 'f2'}, + ] + + def test_positional_and_dict_binding(self): + def fn(a, b): + return a + b + + assert verify_args(fn, [(1, 2)], None) == [{'a': 1, 'b': 2}] + assert verify_args(fn, [{'a': 1, 'b': 2, 'extra': 3}], None) == [ + {'a': 1, 'b': 2, 'extra': 3} + ] + + def test_dict_missing_required_name_raises(self): + def fn(a, b): + return a + b + + with pytest.raises(ValueError, match='Check the args names'): + verify_args(fn, [{'a': 1}], None) + + def test_var_keyword_allows_arbitrary_dicts(self): + def fn(**kwargs): + return kwargs + + assert verify_args(fn, [{'x': 1}], None) == [{'x': 1}] + + +class TestMiscUtils: + + def test_iterchunks(self): + assert list(iterchunks([1, 2, 3, 4, 5], 2)) == [[1, 2], [3, 4], [5]] + assert list(iterchunks([], 3)) == [] + + def test_agg_data(self): + blob, ranges = agg_data([b'ab', b'cde']) + assert blob == b'abcde' + assert ranges == [(0, 1), (2, 4)] + + def test_split_object_url(self): + assert split_object_url('cos://bucket/dir/file.txt') == ( + 'ibm_cos', 'bucket', 'dir', 'file.txt' + ) + assert split_object_url('s3://bucket/prefix/') == ( + 'aws_s3', 'bucket', 'prefix', '' + ) + assert split_object_url('bucket') == (None, 'bucket', '', '') + assert split_object_url('bucket/key') == (None, 'bucket', '', 'key') + + def test_split_path(self): + assert split_path('/bucket/dir/key') == ('bucket', 'dir/key') + assert split_path('bucket') == ('bucket', None) + assert split_path('bucket/') == ('bucket', '') + + def test_convert_bools_to_string_mutates_in_place(self): + env = {'flag': True, 'count': 1, 'name': 'x'} + assert convert_bools_to_string(env) is env + assert env == {'flag': 'True', 'count': 1, 'name': 'x'} + + def test_is_lithops_worker(self, monkeypatch): + monkeypatch.delenv('LITHOPS_WORKER', raising=False) + assert is_lithops_worker() is False + monkeypatch.setenv('LITHOPS_WORKER', '1') + assert is_lithops_worker() is True + + def test_version_str_and_current_py_version(self): + assert version_str((3, 12, 1)) == '3.12' + assert CURRENT_PY_VERSION == version_str(__import__('sys').version_info) + + def test_verify_runtime_name(self): + verify_runtime_name('python:3.12') + with pytest.raises(AssertionError, match='not valid'): + verify_runtime_name('bad name') + + def test_timeout_handler_raises(self): + with pytest.raises(TimeoutError, match='too slow'): + timeout_handler('too slow', None, None) + + def test_b64_dict_roundtrip(self): + payload = {'a': 1, 'b': 'x'} + assert b64str_to_dict(dict_to_b64str(payload)) == payload + + def test_is_object_processing_function(self): + assert is_object_processing_function(lambda obj: obj) + assert not is_object_processing_function(lambda x: x) + + def test_countdown_latch(self): + latch = CountDownLatch(2) + assert latch.done is False + latch.unlock() + assert latch.done is False + latch.unlock() + assert latch.done is True + latch.wait() + + def test_countdown_latch_wait_returns_immediately_when_already_done(self): + latch = CountDownLatch(0) + latch.wait() + assert latch.done is True + + def test_log_prefix_builds_executor_job_and_call_identity(self): + assert log_prefix('sess-0') == 'ExecutorID sess-0' + assert log_prefix('sess-0', 'M000') == 'ExecutorID sess-0 | JobID M000' + assert log_prefix('sess-0', 'M000', '00007') == ( + 'ExecutorID sess-0 | JobID M000 | CallID 00007' + ) + + def test_log_prefix_omits_job_when_only_call_is_set(self): + # call_id without job_id is unusual but must not invent a JobID segment. + assert log_prefix('sess-0', call_id='00007') == 'ExecutorID sess-0 | CallID 00007' + + def test_as_future_list_and_future_id(self): + future = type('F', (), {'executor_id': 'e', 'job_id': 'j', 'call_id': 'c'})() + assert _as_future_list(future) == [future] + plain = [future] + assert _as_future_list(plain) is plain + assert _future_id(future) == ('e', 'j', 'c') + + def test_create_executor_id_reuses_session_and_increments(self, monkeypatch): + monkeypatch.delenv('__LITHOPS_SESSION_ID', raising=False) + monkeypatch.delenv('__LITHOPS_TOTAL_EXECUTORS', raising=False) + first = create_executor_id(lenght=4) + second = create_executor_id(lenght=4) + session, num = first.rsplit('-', 1) + assert len(session) == 4 + assert num == '0' + assert second == f'{session}-1' + assert get_executor_id() == second + + def test_monitoring_queue_chain_matches_the_shapes_in_use(self, monkeypatch): + # These are the chains the id-derived formula produced, and every id + # shape Lithops builds today. The chain must not change for them + monkeypatch.delenv(MONITORING_QUEUES_ENV, raising=False) + assert monitoring_queues('sess-0') == ['lithops-sess-0'] + + # An executor created inside a worker task, whose session id is + # job_key-call_id, or inside a remote invoker, whose session id is + # job_key: both inherit the client's queue through the environment + monkeypatch.setenv(MONITORING_QUEUES_ENV, '["lithops-sess-0"]') + assert monitoring_queues('sess-0-M000-00000-0') == [ + 'lithops-sess-0', 'lithops-sess-0-M000-00000-0' + ] + assert monitoring_queues('sess-0-M000-0') == [ + 'lithops-sess-0', 'lithops-sess-0-M000-0' + ] + + def test_monitoring_queue_chain_goes_deeper_than_the_old_formula( + self, monkeypatch + ): + # The id-derived formula emitted 'lithops-sess-0-M000-0-M000' here, + # a session id no monitor ever declares a queue for + monkeypatch.setenv( + MONITORING_QUEUES_ENV, + '["lithops-sess-0", "lithops-sess-0-M000-0"]', + ) + assert monitoring_queues('sess-0-M000-0-M000-0') == [ + 'lithops-sess-0', + 'lithops-sess-0-M000-0', + 'lithops-sess-0-M000-0-M000-0', + ] + + def test_monitoring_queues_does_not_repeat_a_queue(self, monkeypatch): + monkeypatch.setenv(MONITORING_QUEUES_ENV, '["lithops-sess-0"]') + assert monitoring_queues('sess-0') == ['lithops-sess-0'] + + def test_monitoring_queues_ignores_a_malformed_environment( + self, monkeypatch + ): + monkeypatch.setenv(MONITORING_QUEUES_ENV, 'not json') + assert monitoring_queues('sess-0') == ['lithops-sess-0'] + + def test_monitoring_queue_name(self): + assert monitoring_queue_name('sess-0') == 'lithops-sess-0' + + def test_sizeof_fmt(self): + assert sizeof_fmt(0) == '0.0B' + assert sizeof_fmt(500) == '500.0B' + assert sizeof_fmt(2048).endswith('KiB') + assert sizeof_fmt(-2048).startswith('-') + + def test_bytes_b64_roundtrip(self): + payload = b'hello' + assert b64str_to_bytes(bytes_to_b64str(payload)) == payload + + def test_sdb_to_dict(self): + item = {'Attributes': [{'Name': 'a', 'Value': '1'}, {'Name': 'b', 'Value': 'x'}]} + assert sdb_to_dict(item) == {'a': '1', 'b': 'x'} + + def test_is_unix_and_linux(self, monkeypatch): + monkeypatch.setattr('lithops.utils.platform.system', lambda: 'Darwin') + assert is_unix_system() is True + assert is_linux_system() is False + monkeypatch.setattr('lithops.utils.platform.system', lambda: 'Windows') + assert is_unix_system() is False + monkeypatch.setattr('lithops.utils.platform.system', lambda: 'Linux') + assert is_linux_system() is True + + def test_is_notebook_without_ipython(self): + assert is_notebook() is False + + def test_create_futures_list_attaches_executor(self): + executor = type('E', (), {'config': {'x': 1}})() + fl = create_futures_list(['a'], executor) + assert isinstance(fl, FuturesList) + assert list(fl) == ['a'] + assert fl.executor is executor + assert fl.config == {'x': 1} + + def test_split_object_url_unknown_scheme_is_kept(self): + assert split_object_url('gs://bucket/dir/file') == ( + 'gs', 'bucket', 'dir', 'file' + ) + + def test_format_data_without_extra_args_returns_same_list(self): + data = [1, 2] + assert format_data(data, None) is data + assert format_data(data, []) is data + + def test_verify_args_with_tuple_extra_args(self): + def fn(a, b): + return a + b + + assert verify_args(fn, [(1,)], (2,)) == [{'a': 1, 'b': 2}] + + def test_get_docker_path_prefers_docker_then_podman(self, monkeypatch): + monkeypatch.setattr('lithops.utils.shutil.which', lambda name: { + 'docker': '/usr/bin/docker', + 'podman': None, + }[name]) + assert get_docker_path() == '/usr/bin/docker' + monkeypatch.setattr('lithops.utils.shutil.which', lambda name: { + 'docker': None, + 'podman': '/usr/bin/podman', + }[name]) + assert get_docker_path() == '/usr/bin/podman' + monkeypatch.setattr('lithops.utils.shutil.which', lambda name: None) + with pytest.raises(Exception, match='docker/podman command not found'): + get_docker_path() + + def test_docker_login_requires_credentials(self): + with pytest.raises(Exception, match='docker_user and docker_password'): + docker_login(None, None, 'docker.io') + with pytest.raises(Exception, match='docker_server is required'): + docker_login('u', 'p', '') + + def test_get_default_container_name_variants(self): + cfg = {'docker_server': 'docker.io', 'docker_user': 'alice'} + name = get_default_container_name('ibm_cf', cfg, 'lithops') + assert name.startswith('docker.io/alice/lithops-v') + cfg = {'docker_server': 'icr.io', 'docker_namespace': 'ns'} + name = get_default_container_name('ibm_cf', cfg, 'lithops') + assert name.startswith('icr.io/ns/lithops-v') + cfg = { + 'docker_server': 'us-docker.pkg.dev', + 'region': 'us', + 'project_name': 'proj', + } + name = get_default_container_name('gcp', cfg, 'lithops') + assert name.startswith('us-docker.pkg.dev/proj/lithops/lithops-v') + cfg = {'docker_server': 'example.registry'} + name = get_default_container_name('k8s', cfg, 'lithops') + assert name.startswith('example.registry/lithops-v') + + def test_get_default_container_name_missing_docker_user(self): + with pytest.raises(Exception, match='docker_user'): + get_default_container_name('ibm_cf', {'docker_server': 'docker.io'}, 'r') + + def test_is_podman(self, monkeypatch): + monkeypatch.setattr( + 'lithops.utils.sp.check_output', lambda *a, **k: b'podman' + ) + assert is_podman('/usr/bin/podman') is True + monkeypatch.setattr( + 'lithops.utils.sp.check_output', + lambda *a, **k: (_ for _ in ()).throw(Exception('nope')), + ) + assert is_podman('/usr/bin/docker') is False + + def test_create_handler_zip(self, tmp_path): + from lithops.utils import create_handler_zip + entry = tmp_path / 'entry.py' + entry.write_text('print(1)\n') + dest = tmp_path / 'handler.zip' + create_handler_zip(str(dest), [str(entry)]) + assert zipfile.is_zipfile(dest) + with zipfile.ZipFile(dest) as zf: + names = zf.namelist() + assert 'entry.py' in names + assert any(name.startswith('lithops/') for name in names) + + def test_create_handler_zip_skips_output_zip_and_caches( + self, tmp_path, monkeypatch + ): + from lithops.utils import create_handler_zip + + pkg = tmp_path / 'lithops' + pkg.mkdir() + (pkg / '__init__.py').write_text('') + (pkg / '__pycache__').mkdir() + (pkg / '__pycache__' / 'mod.cpython-312.pyc').write_bytes(b'nope') + pytest_cache = pkg / '.pytest_cache' + pytest_cache.mkdir() + (pytest_cache / 'v').write_text('nope') + leftover = pkg / 'leftover.zip' + leftover.write_bytes(b'PK' + b'\x00' * 32) + + monkeypatch.setattr('lithops.__file__', str(pkg / '__init__.py')) + + entry = tmp_path / 'entry.py' + entry.write_text('print(1)\n') + dest = pkg / 'handler.zip' + create_handler_zip(str(dest), [str(entry)]) + + with zipfile.ZipFile(dest) as zf: + names = zf.namelist() + assert 'entry.py' in names + assert 'lithops/__init__.py' in names + assert not any('__pycache__' in name for name in names) + assert not any('.pytest_cache' in name for name in names) + assert not any(name.endswith('.zip') for name in names) + + def test_wrapped_streaming_body_read_seek_and_eof(self): + body = WrappedStreamingBody(io.BytesIO(b'hello world'), 11) + assert body.read(5) == b'hello' + assert body.tell() == 5 + # Historical: whence=0 does not apply offset; it returns the current pos. + assert body.seek(0) == 5 + assert body.seek(2, 1) == 7 + assert body.seek(0, 2) == 11 + with pytest.raises(Exception, match='Unsupported'): + body.seek(-1, 2) + + class Empty: + def read(self, n=None): + return "" + + with pytest.raises(EOFError): + WrappedStreamingBody(Empty(), 1).read() + + def test_wrapped_streaming_body_partition_first_chunk(self): + data = b'aaa\nbbb\nccc\n' + part = WrappedStreamingBodyPartition( + io.BytesIO(data), size=len(data), byterange=(0, len(data) - 1) + ) + assert b'aaa' in part.read(100) + + def test_iterchunks_chunk_larger_than_list(self): + assert list(iterchunks([1, 2], 10)) == [[1, 2]] + + def test_agg_data_empty(self): + blob, ranges = agg_data([]) + assert blob == b'' + assert ranges == [] + + def test_setup_logger_none_is_noop(self): + with patch('lithops.utils.logging.config.dictConfig') as cfg: + setup_lithops_logger(None) + setup_lithops_logger('none') + setup_lithops_logger('NONE') + cfg.assert_not_called() + + def test_setup_logger_debug_uses_console_handler(self): + with patch('lithops.utils.logging.config.dictConfig') as cfg: + setup_lithops_logger('debug') + config = cfg.call_args[0][0] + assert config['loggers']['lithops']['handlers'] == ['console_handler'] + assert config['handlers']['console_handler']['level'] == logging.DEBUG + assert config['handlers']['console_handler']['class'] == ( + 'lithops.utils.ShutdownSafeStreamHandler' + ) + + def test_shutdown_safe_handler_ignores_closed_stream(self, capsys): + stream = io.StringIO() + handler = ShutdownSafeStreamHandler(stream) + test_logger = logging.getLogger('lithops.test_closed_stream') + test_logger.handlers = [handler] + test_logger.propagate = False + test_logger.setLevel(logging.DEBUG) + stream.close() + test_logger.debug('after close') + + class ClosedWrite: + closed = False + + def write(self, msg): + raise ValueError('I/O operation on closed file.') + + def flush(self): + return None + + handler.stream = ClosedWrite() + test_logger.debug('after failed write') + captured = capsys.readouterr() + assert 'Logging error' not in captured.err + assert 'Logging error' not in captured.out + + def test_setup_logger_filename_uses_file_handler(self, tmp_path): + log_file = str(tmp_path / 'lithops.log') + with patch('lithops.utils.logging.config.dictConfig') as cfg: + setup_lithops_logger('info', filename=log_file) + config = cfg.call_args[0][0] + assert config['loggers']['lithops']['handlers'] == ['file_handler'] + assert config['handlers']['file_handler']['filename'] == log_file + + def test_run_command_check_call_by_default(self): + with patch('lithops.utils.sp.check_call') as call: + run_command('echo hello') + call.assert_called_once() + assert call.call_args[0][0] == ['echo', 'hello'] + + def test_run_command_return_result_strips_quotes(self): + with patch('lithops.utils.sp.check_output', return_value=' "ok" '): + assert run_command('echo x', return_result=True) == 'ok' + + def test_run_command_input_uses_check_output_bytes(self): + with patch('lithops.utils.sp.check_output', return_value=b'') as out: + run_command('cat', input='secret') + assert out.call_args.kwargs['input'] == b'secret' + + def test_find_free_port_returns_int(self): + first = find_free_port() + second = find_free_port() + assert isinstance(first, int) and isinstance(second, int) + assert 0 < first < 65536 + assert 0 < second < 65536 + + def test_futures_list_map_wait_get_result_and_pickle(self): + class Item: + def __init__(self): + self._produce_output = True + + executor = MagicMock() + mapped = [Item()] + executor.map.return_value = mapped + executor.wait.return_value = ([], []) + executor.get_result.return_value = [1] + + fl = FuturesList([Item()]) + fl.executor = executor + fl.config = {} + result = fl.map(lambda x: x, sync=True) + executor.wait.assert_called() + executor.map.assert_called() + assert result is fl + assert list(fl) == mapped + fl.wait() + assert fl.get_result() == [1] + + fl2 = FuturesList([1, 2]) + fl2.executor = object() + dumped = pickle.dumps(fl2) + assert fl2.executor is None + loaded = pickle.loads(dumped) + assert list(loaded) == [1, 2] + assert loaded.executor is None + + def test_wrapped_streaming_body_partition_middle_chunk_discards_partial_row(self): + data = b'aaa\nbbb\nccc\n' + part = WrappedStreamingBodyPartition( + io.BytesIO(data[2:]), size=len(data[2:]), byterange=(2, len(data) - 1) + ) + assert part.read(100) == b'bbb\nccc\n' + + def test_wrapped_streaming_body_partition_readline_discards_partial_row(self): + class Stream: + def __init__(self, payload): + self._raw_stream = io.BytesIO(payload) + + def read(self, n=None): + return self._raw_stream.read(-1 if n is None else n) + + data = b'aaa\nbbb\nccc\n' + part = WrappedStreamingBodyPartition( + Stream(data[2:]), size=len(data[2:]), byterange=(2, len(data) - 1) + ) + assert part.readline() == b'bbb\n' + + def test_create_handler_zip_removes_partial_zip_on_write_error(self, tmp_path): + from lithops.utils import create_handler_zip + + entry = tmp_path / 'entry.py' + entry.write_text('print(1)\n') + dest = tmp_path / 'handler.zip' + + def boom(self, *args, **kwargs): + raise RuntimeError('disk full') + + with patch.object(zipfile.ZipFile, 'write', boom): + with pytest.raises(Exception, match='Unable to create'): + create_handler_zip(str(dest), [str(entry)]) + assert not dest.exists() + + def test_create_handler_zip_removes_partial_zip_on_keyboard_interrupt( + self, tmp_path + ): + from lithops.utils import create_handler_zip + + entry = tmp_path / 'entry.py' + entry.write_text('print(1)\n') + dest = tmp_path / 'handler.zip' + + def boom(self, *args, **kwargs): + raise KeyboardInterrupt() + + with patch.object(zipfile.ZipFile, 'write', boom): + with pytest.raises(KeyboardInterrupt): + create_handler_zip(str(dest), [str(entry)]) + assert not dest.exists() diff --git a/lithops/tests/test_wait.py b/lithops/tests/test_wait.py new file mode 100644 index 000000000..561be768b --- /dev/null +++ b/lithops/tests/test_wait.py @@ -0,0 +1,399 @@ +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import signal +import threading +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from lithops.utils import is_unix_system + +from lithops.utils import FuturesList +from lithops.wait import ( + ALL_COMPLETED, + ALWAYS, + ANY_COMPLETED, + WAIT_DUR_SEC, + _as_future_list, + _check_done, + _create_executors_data_from_futures, + _future_is_complete, + _get_executor_data, + _partition_futures, + _poll_sleep_sec, + _ready_futures, + get_result, + wait, +) + + +class FakeFuture: + def __init__(self, *, done=False, success=False, ready=False, executor_id='sess-0', + job_id='M000', call_id='00000', storage_backend='localhost', + result=None, produce_output=True, futures=False): + self.done = done + self.success = success + self.ready = ready + self.executor_id = executor_id + self.job_id = job_id + self.call_id = call_id + self._storage_config = {'backend': storage_backend} + self._result = result + self._produce_output = produce_output + self.futures = futures + self._new_futures = None + + def result(self, throw_except=True, internal_storage=None): + self.done = True + return self._result + + def status(self, throw_except=True, internal_storage=None): + self.success = True + return {'ok': True} + + +class TestWaitHelpers: + + def test_as_future_list_wraps_single_future(self): + future = FakeFuture() + assert _as_future_list(future) == [future] + + def test_as_future_list_keeps_list_and_futures_list(self): + plain = [FakeFuture()] + futures_list = FuturesList(plain) + assert _as_future_list(plain) is plain + assert _as_future_list(futures_list) is futures_list + + def test_future_is_complete_depends_on_download_results(self): + success_only = FakeFuture(done=False, success=True) + assert _future_is_complete(success_only, download_results=False) is True + assert _future_is_complete(success_only, download_results=True) is False + + finished = FakeFuture(done=True, success=True) + assert _future_is_complete(finished, download_results=True) is True + + def test_partition_preserves_order(self): + first = FakeFuture(done=True, success=True) + second = FakeFuture() + third = FakeFuture(done=False, success=True) + + done, not_done = _partition_futures( + [first, second, third], download_results=False + ) + assert done == [first, third] + assert not_done == [second] + + done, not_done = _partition_futures( + [first, second, third], download_results=True + ) + assert done == [first] + assert not_done == [second, third] + + def test_partition_empty(self): + assert _partition_futures([], False) == ([], []) + + def test_check_done_any_completed(self): + pending = FakeFuture() + finished = FakeFuture(done=True, success=True) + assert _check_done([pending, pending], ANY_COMPLETED, False) is False + assert _check_done([pending, finished], ANY_COMPLETED, False) is True + + def test_check_done_percentage_and_all_completed(self): + finished = FakeFuture(done=True, success=True) + pending = FakeFuture() + fs = [finished, pending] + assert _check_done(fs, 50, False) is True + assert _check_done(fs, ALL_COMPLETED, False) is False + assert _check_done([finished, finished], ALL_COMPLETED, False) is True + + def test_check_done_always_is_immediately_true(self): + assert _check_done([FakeFuture()], ALWAYS, False) is True + + +class TestWait: + + def test_empty_input_returns_two_empty_lists(self): + assert wait([]) == ([], []) + assert wait(None) == ([], []) + + def test_returns_immediately_when_all_complete(self): + future = FakeFuture(done=True, success=True) + done, not_done = wait([future], show_progressbar=False) + assert done == [future] + assert not_done == [] + + def test_wraps_single_complete_future(self): + future = FakeFuture(done=True, success=True) + done, not_done = wait(future, show_progressbar=False) + assert done == [future] + assert not_done == [] + + +class TestCreateExecutorsData: + + @patch('lithops.wait.InternalStorage') + def test_groups_futures_and_reuses_matching_storage(self, mock_storage_cls): + internal = MagicMock() + internal.backend = 'localhost' + first = FakeFuture(executor_id='a') + second = FakeFuture(executor_id='a', call_id='00001') + third = FakeFuture(executor_id='b', storage_backend='s3') + + groups = _create_executors_data_from_futures( + [first, second, third], internal + ) + by_id = {group.executor_id: group for group in groups} + + assert set(by_id) == {'a', 'b'} + assert by_id['a'].futures == [first, second] + assert by_id['a'].internal_storage is internal + assert by_id['b'].futures == [third] + mock_storage_cls.assert_called_once_with(third._storage_config) + assert by_id['b'].internal_storage is mock_storage_cls.return_value + + +class TestPollSleep: + + def test_localhost_and_non_storage_use_short_interval(self): + local = SimpleNamespace(type='storage', storage_backend='localhost') + rabbit = SimpleNamespace(type='rabbitmq', storage_backend='s3') + assert _poll_sleep_sec(local, None) == 0.1 + assert _poll_sleep_sec(rabbit, 3) == 0.1 + + def test_remote_storage_uses_wait_dur_or_default(self): + remote = SimpleNamespace(type='storage', storage_backend='s3') + assert _poll_sleep_sec(remote, None) == WAIT_DUR_SEC + assert _poll_sleep_sec(remote, 3) == 3 + # Historical: 0 is falsy so the default interval is used. + assert _poll_sleep_sec(remote, 0) == WAIT_DUR_SEC + + +class TestReadyFuturesAndExecutorData: + + def test_ready_futures_status_only_includes_ready_and_pending(self): + ready = FakeFuture(ready=True, call_id='00000') + success = FakeFuture(success=True, call_id='00001') + pending = FakeFuture(call_id='00002') + exec_data = SimpleNamespace(futures=[ready, success, pending]) + assert _ready_futures(exec_data, download_results=False) == [ready] + + def test_ready_futures_download_includes_success_until_done(self): + success = FakeFuture(ready=True, success=True, call_id='00000') + done = FakeFuture(done=True, success=True, ready=True, call_id='00001') + exec_data = SimpleNamespace(futures=[success, done]) + assert _ready_futures(exec_data, download_results=True) == [success] + + def test_get_executor_data_fetches_status_and_extends_new_futures(self): + parent = FakeFuture(ready=True, call_id='00000') + child = FakeFuture(call_id='00001') + parent._new_futures = [child] + exec_data = SimpleNamespace(futures=[parent], internal_storage=MagicMock()) + fs = [parent] + pbar = MagicMock() + pbar.n = 0 + pbar.total = 1 + + fetched = _get_executor_data( + fs, exec_data, download_results=False, throw_except=True, + threadpool_size=2, pbar=pbar, + ) + + assert fetched == 1 + assert parent.success is True + assert child in fs + assert child in exec_data.futures + assert pbar.total == 2 + pbar.update.assert_called() + + def test_get_executor_data_downloads_results(self): + future = FakeFuture(ready=True, success=True, result=42) + exec_data = SimpleNamespace(futures=[future], internal_storage=MagicMock()) + fetched = _get_executor_data( + [future], exec_data, download_results=True, throw_except=True, + threadpool_size=1, pbar=None, + ) + assert fetched == 1 + assert future.done is True + + +class TestWaitPolling: + + def test_always_polls_once_without_looping(self): + future = FakeFuture() + monitor = MagicMock() + monitor.type = 'storage' + monitor.storage_backend = 'localhost' + internal = MagicMock() + internal.backend = 'localhost' + with patch('lithops.wait._get_executor_data', return_value=0) as get: + wait( + [future], + return_when=ALWAYS, + show_progressbar=False, + job_monitor=monitor, + internal_storage=internal, + ) + get.assert_called_once() + + def test_keyboard_interrupt_reraises_after_logging(self): + future = FakeFuture() + with patch( + 'lithops.wait._create_executors_data_from_futures', + side_effect=KeyboardInterrupt, + ): + with pytest.raises(KeyboardInterrupt): + wait([future], show_progressbar=False) + + def test_starts_and_stops_monitor_when_none_provided(self): + future = FakeFuture() + monitor = MagicMock() + monitor.type = 'storage' + monitor.storage_backend = 'localhost' + monitor.is_alive.return_value = True + internal = MagicMock() + internal.backend = 'localhost' + + def get_data(fs, exec_data, **kwargs): + future.success = True + return 1 + + with patch('lithops.wait.JobMonitor', return_value=monitor) as cls, \ + patch('lithops.wait._get_executor_data', side_effect=get_data), \ + patch('lithops.wait.time.sleep'): + wait( + [future], + show_progressbar=False, + internal_storage=internal, + ) + cls.assert_called_once() + monitor.start.assert_called_once() + monitor.stop.assert_called_once() + + def test_timeout_registers_sigalrm_with_interpolated_message(self): + if not is_unix_system(): + pytest.skip('SIGALRM waiting timeout is unix-only') + future = FakeFuture() + monitor = MagicMock() + monitor.type = 'storage' + monitor.storage_backend = 'localhost' + monitor.is_alive.return_value = True + internal = MagicMock() + internal.backend = 'localhost' + handlers = {} + + def fake_signal(sig, handler): + handlers[sig] = handler + + def get_data(fs, exec_data, **kwargs): + future.success = True + return 1 + + with patch('lithops.wait.signal.signal', side_effect=fake_signal), \ + patch('lithops.wait.signal.alarm') as alarm, \ + patch('lithops.wait._get_executor_data', side_effect=get_data), \ + patch('lithops.wait.time.sleep'): + wait( + [future], + timeout=17, + show_progressbar=False, + job_monitor=monitor, + internal_storage=internal, + ) + + alarm.assert_any_call(17) + alarm.assert_called_with(0) + assert 'Timeout of 17 seconds exceeded' in handlers[signal.SIGALRM].args[0] + + def test_all_completed_restarts_dead_monitor_and_sleeps_on_empty_poll(self): + future = FakeFuture() + monitor = MagicMock() + monitor.type = 'storage' + monitor.storage_backend = 'localhost' + monitor.is_alive.side_effect = [False, True] + internal = MagicMock() + internal.backend = 'localhost' + polls = {'n': 0} + + def get_data(fs, exec_data, **kwargs): + polls['n'] += 1 + if polls['n'] == 1: + return 0 + future.success = True + future.done = True + return 3 + + sleeps = [] + test_thread = threading.current_thread() + + def sleep(seconds): + if threading.current_thread() is test_thread: + sleeps.append(seconds) + + with patch('lithops.wait._get_executor_data', side_effect=get_data), \ + patch('lithops.wait.time.sleep', side_effect=sleep): + wait( + [future], + return_when=ALL_COMPLETED, + show_progressbar=False, + job_monitor=monitor, + internal_storage=internal, + ) + + monitor.start.assert_called_once_with(fs=[future]) + assert sleeps == [0.1, 0] + + def test_wait_tracks_nested_futures_until_they_complete(self): + child = FakeFuture(call_id='00001') + parent = FakeFuture(ready=True, call_id='00000') + + def parent_status(**kwargs): + parent.success = True + parent._new_futures = [child] + child.ready = True + return {'ok': True} + + parent.status = parent_status + monitor = MagicMock() + monitor.type = 'storage' + monitor.storage_backend = 'localhost' + monitor.is_alive.return_value = True + internal = MagicMock() + internal.backend = 'localhost' + + with patch('lithops.wait.time.sleep'): + done, not_done = wait( + [parent], + show_progressbar=False, + job_monitor=monitor, + internal_storage=internal, + ) + + assert parent.success is True + assert child.success is True + assert child in done + assert parent in done + assert not_done == [] + + +class TestGetResult: + + def test_returns_produced_results_and_skips_nested(self): + produced = FakeFuture(done=True, success=True, result=1) + nested = FakeFuture(done=True, success=True, result=2, futures=True) + silent = FakeFuture(done=True, success=True, result=3, produce_output=False) + assert get_result( + [produced, nested, silent], show_progressbar=False + ) == [1] diff --git a/lithops/tests/test_worker.py b/lithops/tests/test_worker.py new file mode 100644 index 000000000..4912c4749 --- /dev/null +++ b/lithops/tests/test_worker.py @@ -0,0 +1,1169 @@ +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import io +import os +import pickle +import sys +import threading +from queue import Empty, Queue +from types import SimpleNamespace +from unittest.mock import MagicMock, mock_open, patch + +import pytest + +import lithops.worker.handler as handler_module +from lithops.constants import JOBS_PREFIX, MODULES_DIR +from lithops.storage.utils import CloudObject, CloudObjectLocal, CloudObjectUrl +from lithops.utils import bytes_to_b64str, is_unix_system +from lithops.worker import function_handler, function_invoker +from lithops.worker.handler import ( + ShutdownSentinel, + TaskJar, + create_job, + prepare_and_run_task, + task_consumer, + run_task, +) +from lithops.worker.jobrunner import JobRunner, JobStats, _prepare_args +from lithops.worker.status import ( + CallStatus, + RabbitmqCallStatus, + StorageCallStatus, + create_call_status, +) +from lithops.worker.utils import ( + LogStream, + SystemMonitor, + custom_redirection, + free_disk_space, + get_function_and_modules, + get_function_data, + get_memory_usage, + get_runtime_metadata, + memory_monitor_worker, + peak_memory, + psutil_found, +) + + +def _job_config(monitoring='storage', **lithops): + return { + 'lithops': { + 'storage': 'localhost', + 'backend': 'localhost', + 'monitoring': monitoring, + **lithops, + }, + 'localhost': {'storage_bucket': 'bucket'}, + 'rabbitmq': {'amqp_url': 'amqp://guest:guest@localhost:5672'}, + } + + +def _echo(x): + return x + + +def _add(x): + return x + 1 + + +def _big(x): + return 'y' * 9000 + + +def _none(x): + return None + + +def _boom(x): + raise ValueError('nope') + + +def _obj_fn(obj): + return 1 + + +def _reduce_fn(results): + return results + + +def _with_id_and_storage(x, id, storage): + return x + + +class _Adder: + def __call__(self, x): + return x + + +def _return_futures_list(x): + from lithops.utils import FuturesList + return FuturesList() + + +def _record_task(task): + """Leave a file behind, as a forked worker cannot share state in memory""" + name = f'{task.call_id}-{task.data.decode()}' + with open(os.path.join(task.out_dir, name), 'w') as fid: + fid.write(str(os.getpid())) + + +def _task(**kwargs): + values = dict( + extra_env={}, + config=_job_config(), + job_key='ek-j0', + call_id='00000', + log_level='ERROR', + runtime_name='rt', + runtime_memory=None, + execution_timeout=10, + start_tstamp=1.0, + host_submit_tstamp=0.5, + job_id='j0', + executor_id='ek', + chunksize=1, + func=pickle.dumps(_echo), + data=pickle.dumps({'x': 1}), + stats_file='stats.txt', + ) + values.update(kwargs) + return SimpleNamespace(**values) + + +class TestPackageExports: + + def test_exports(self): + import lithops.worker as worker + assert worker.function_handler is function_handler + assert worker.function_invoker is function_invoker + assert worker.__all__ == ['function_handler', 'function_invoker'] + + +class TestCreateJob: + + def test_loads_func_and_data(self): + payload = {'config': _job_config(), 'job_key': 'jk'} + with patch( + 'lithops.worker.handler.extract_storage_config', return_value={} + ): + with patch('lithops.worker.handler.InternalStorage') as store: + with patch( + 'lithops.worker.handler.get_function_and_modules', + return_value=b'f', + ) as gf: + with patch( + 'lithops.worker.handler.get_function_data', + return_value=[b'd'], + ) as gd: + job = create_job(payload) + assert job.func == b'f' + assert job.data == [b'd'] + gf.assert_called_once() + gd.assert_called_once() + store.assert_called_once() + + +class TestTaskConsumer: + + def test_runs_tasks_then_stops_on_sentinel(self): + q = Queue() + task = _task() + q.put((task, '00001', b'data')) + q.put(ShutdownSentinel()) + init = MagicMock() + cb = MagicMock() + with patch('lithops.worker.handler.prepare_and_run_task') as run: + task_consumer(3, q, initializer=init, callback=cb) + run.assert_called_once_with(task) + assert task.call_id == '00001' + assert task.data == b'data' + init.assert_called_once_with(3, task) + cb.assert_called_once_with(3, task) + + def test_none_initializer_and_callback_are_skipped(self): + q = Queue() + q.put(ShutdownSentinel()) + task_consumer(0, q) + + def test_empty_and_broken_pipe_stop_the_loop(self): + q = MagicMock() + q.get.side_effect = Empty() + task_consumer(0, q) + q.get.side_effect = BrokenPipeError() + task_consumer(0, q) + + def test_a_failed_task_does_not_stop_the_worker(self): + q = Queue() + q.put((_task(), '00001', b'a')) + q.put((_task(), '00002', b'b')) + q.put(ShutdownSentinel()) + with patch( + 'lithops.worker.handler.prepare_and_run_task', + side_effect=[OSError('no space left'), None], + ) as run: + with patch.object(handler_module.logger, 'error') as log_error: + task_consumer(0, q) + assert run.call_count == 2 + assert 'failed to run task 00001' in log_error.call_args[0][0] + + +class TestTaskJar: + + def _jar(self, n_calls): + job = _task( + call_ids=[f'{i:05}' for i in range(n_calls)], + data=[f'd{i}'.encode() for i in range(n_calls)], + ) + return TaskJar(job) + + def test_dispatch_hands_out_every_call_in_order(self): + jar = self._jar(3) + jar.dispatch() + assert [jar.get()[1:] for _ in range(3)] == [ + ('00000', b'd0'), ('00001', b'd1'), ('00002', b'd2') + ] + with pytest.raises(Empty): + jar.get() + jar.close_reader() + + def test_get_returns_the_job_and_survives_data_rebinding(self): + jar = self._jar(2) + jar.dispatch() + task, call_id, data = jar.get() + assert task is jar.job + # task_consumer rebinds job.data to the running task + task.data = data + assert jar.get()[1:] == ('00001', b'd1') + jar.close_reader() + + def test_dispatch_survives_workers_that_died(self): + jar = self._jar(2) + jar.close_reader() + with patch.object(handler_module.logger, 'error') as log_error: + jar.dispatch() + assert 'exited before consuming all tasks' in log_error.call_args[0][0] + + def test_tokens_are_never_split_by_a_partial_write(self): + # More calls than a pipe buffer holds, so os.write returns short + jar = self._jar(40000) + writer = threading.Thread(target=jar.dispatch) + writer.start() + try: + claimed = [jar.get()[1] for _ in range(40000)] + finally: + writer.join() + assert claimed == [f'{i:05}' for i in range(40000)] + jar.close_reader() + + +class TestFunctionHandler: + + def test_single_worker_uses_threading_queue(self): + job = _task(worker_processes=4, call_ids=['00000'], data=[b'd']) + with patch('lithops.worker.handler.create_job', return_value=job): + with patch('lithops.worker.handler.setup_lithops_logger'): + with patch( + 'lithops.worker.handler.task_consumer' + ) as consumer: + function_handler({}) + consumer.assert_called_once() + assert consumer.call_args[0][0] == 0 + + def test_multi_worker_starts_processes_and_joins(self): + job = _task( + worker_processes=2, call_ids=['00000', '00001'], data=[b'a', b'b'] + ) + proc = MagicMock() + ctx = MagicMock() + ctx.Process.return_value = proc + with patch('lithops.worker.handler.create_job', return_value=job): + with patch('lithops.worker.handler.setup_lithops_logger'): + with patch('lithops.worker.handler._MP_CTX', ctx): + function_handler({}) + ctx.Manager.assert_not_called() + assert ctx.Process.call_count == 2 + assert proc.start.call_count == 2 + assert proc.join.call_count == 2 + + def test_multi_worker_runs_every_task_once(self): + n_tasks = 6 + job = _task( + worker_processes=3, + call_ids=[f'{i:05}' for i in range(n_tasks)], + data=[f'd{i}'.encode() for i in range(n_tasks)], + ) + seen = [] + with patch('lithops.worker.handler.create_job', return_value=job): + with patch('lithops.worker.handler.setup_lithops_logger'): + with patch( + 'lithops.worker.handler.prepare_and_run_task', + side_effect=lambda task: seen.append( + (task.call_id, task.data) + ), + ): + with patch( + 'lithops.worker.handler._run_process_pool', + new=handler_module._run_thread_pool, + ): + function_handler({}) + assert sorted(seen) == [ + (f'{i:05}', f'd{i}'.encode()) for i in range(n_tasks) + ] + + @pytest.mark.skipif( + not is_unix_system(), reason='the process pool needs fork' + ) + # pytest itself is multi-threaded, which fork warns about since 3.12 + @pytest.mark.filterwarnings('ignore:.*fork.*:DeprecationWarning') + def test_process_pool_runs_every_task_in_a_child(self, tmp_path): + n_tasks = 8 + job = _task( + worker_processes=4, + call_ids=[f'{i:05}' for i in range(n_tasks)], + data=[f'd{i}'.encode() for i in range(n_tasks)], + out_dir=str(tmp_path), + ) + with patch( + 'lithops.worker.handler.prepare_and_run_task', new=_record_task + ): + handler_module._run_process_pool(job, job.worker_processes) + + done = sorted(p.name for p in tmp_path.iterdir()) + assert done == [f'{i:05}-d{i}' for i in range(n_tasks)] + pids = {p.read_text() for p in tmp_path.iterdir()} + assert str(os.getpid()) not in pids + + def test_removes_module_path_and_total_executors(self): + job = _task(worker_processes=1, call_ids=['00000'], data=[b'd']) + module_path = os.path.join(MODULES_DIR, job.job_key) + sys.path.append(module_path) + os.environ['__LITHOPS_TOTAL_EXECUTORS'] = '2' + try: + with patch('lithops.worker.handler.create_job', return_value=job): + with patch('lithops.worker.handler.setup_lithops_logger'): + with patch('lithops.worker.handler.task_consumer'): + function_handler({}) + assert module_path not in sys.path + assert '__LITHOPS_TOTAL_EXECUTORS' not in os.environ + finally: + if module_path in sys.path: + sys.path.remove(module_path) + os.environ.pop('__LITHOPS_TOTAL_EXECUTORS', None) + + +class TestPrepareAndRunTask: + + def test_sets_env_creates_dir_and_clears_extra_env( + self, tmp_path, monkeypatch + ): + monkeypatch.setattr( + 'lithops.worker.handler.LITHOPS_TEMP_DIR', str(tmp_path) + ) + os.environ.pop('__LITHOPS_ACTIVATION_ID', None) + extra = {'FOO': 'bar'} + task = _task(extra_env=extra, job_key='jk', call_id='c1') + with patch('lithops.worker.handler.run_task') as run: + prepare_and_run_task(task) + run.assert_called_once_with(task) + assert os.environ['LITHOPS_WORKER'] == 'True' + assert os.environ['PYTHONUNBUFFERED'] == 'True' + assert 'FOO' not in os.environ + assert os.path.isdir(task.task_dir) + assert task.task_dir == os.path.join( + str(tmp_path), 'bucket', JOBS_PREFIX, 'jk', 'c1' + ) + assert task.log_file == os.path.join(task.task_dir, 'execution.log') + assert task.stats_file == os.path.join(task.task_dir, 'job_stats.txt') + assert len(os.environ['__LITHOPS_ACTIVATION_ID']) == 12 + + def test_keeps_existing_activation_id(self, tmp_path, monkeypatch): + monkeypatch.setattr( + 'lithops.worker.handler.LITHOPS_TEMP_DIR', str(tmp_path) + ) + os.environ['__LITHOPS_ACTIVATION_ID'] = 'alreadythere1' + task = _task(extra_env={}) + with patch('lithops.worker.handler.run_task'): + prepare_and_run_task(task) + assert os.environ['__LITHOPS_ACTIVATION_ID'] == 'alreadythere1' + + +class TestRunTask: + + def _patch_run(self, task, jrp, handler_conn, stats_text=None): + if stats_text is not None: + with open(task.stats_file, 'w') as f: + f.write(stats_text) + status = MagicMock() + cpu = {'usage': [1], 'system': 0.1, 'user': 0.2} + net = {'sent': 3, 'recv': 4} + mem = {'rss': 5, 'vms': 6, 'uss': 7} + monitor = MagicMock() + monitor.get_cpu_info.return_value = cpu + monitor.get_network_io.return_value = net + monitor.get_memory_info.return_value = mem + ctx = MagicMock() + ctx.Pipe.return_value = (handler_conn, MagicMock()) + ctx.Process.return_value = jrp + with patch('lithops.worker.handler.setup_lithops_logger'): + with patch( + 'lithops.worker.handler.extract_storage_config', return_value={} + ): + with patch('lithops.worker.handler.InternalStorage'): + with patch( + 'lithops.worker.handler.create_call_status', + return_value=status, + ): + with patch('lithops.worker.handler._MP_CTX', ctx): + with patch('lithops.worker.handler.JobRunner'): + with patch( + 'lithops.worker.handler.SystemMonitor', + return_value=monitor, + ): + with patch( + 'lithops.worker.handler.is_unix_system', + return_value=True, + ): + run_task(task) + return status + + def test_success_reads_stats_and_sends_events(self, tmp_path): + task = _task() + task.log_stream = MagicMock() + task.log_file = str(tmp_path / 'execution.log') + task.stats_file = str(tmp_path / 'job_stats.txt') + (tmp_path / 'execution.log').write_bytes(b'log') + jrp = MagicMock() + jrp.is_alive.return_value = False + conn = MagicMock() + conn.poll.return_value = True + status = self._patch_run( + task, jrp, conn, 'worker_func_exec_time 1.5\nexception True\n' + ) + status.send_init_event.assert_called_once() + status.send_finish_event.assert_called_once() + added = {c.args[0]: c.args[1] for c in status.add.call_args_list} + assert added['worker_func_exec_time'] == 1.5 + assert added['exception'] is True + assert 'logs' in added + assert 'worker_end_tstamp' in added + task.log_stream.flush.assert_called() + + def test_timeout_raises_handler_timeout_error(self, tmp_path): + task = _task(execution_timeout=7) + task.log_stream = MagicMock() + task.log_file = str(tmp_path / 'execution.log') + task.stats_file = str(tmp_path / 'missing.txt') + jrp = MagicMock() + jrp.is_alive.return_value = True + conn = MagicMock() + status = self._patch_run(task, jrp, conn) + jrp.terminate.assert_called_once() + added = {c.args[0]: c.args[1] for c in status.add.call_args_list} + assert added['exception'] is True + assert 'exc_info' in added + + def test_no_completion_message_is_memory_error(self, tmp_path): + task = _task() + task.log_stream = MagicMock() + task.log_file = str(tmp_path / 'execution.log') + task.stats_file = str(tmp_path / 'missing.txt') + jrp = MagicMock() + jrp.is_alive.return_value = False + conn = MagicMock() + conn.poll.return_value = False + status = self._patch_run(task, jrp, conn) + added = {c.args[0]: c.args[1] for c in status.add.call_args_list} + assert added['exception'] is True + + def test_keyboard_interrupt_skips_finish_event(self, tmp_path): + task = _task() + task.log_stream = MagicMock() + task.log_file = str(tmp_path / 'execution.log') + task.stats_file = str(tmp_path / 'missing.txt') + status = MagicMock() + status.send_init_event.side_effect = KeyboardInterrupt() + with patch('lithops.worker.handler.setup_lithops_logger'): + with patch( + 'lithops.worker.handler.extract_storage_config', return_value={} + ): + with patch('lithops.worker.handler.InternalStorage'): + with patch( + 'lithops.worker.handler.create_call_status', + return_value=status, + ): + run_task(task) + status.send_finish_event.assert_not_called() + + def test_runtime_memory_log_branch(self, tmp_path): + task = _task(runtime_memory=256) + task.log_stream = MagicMock() + task.log_file = str(tmp_path / 'execution.log') + task.stats_file = str(tmp_path / 'missing.txt') + jrp = MagicMock() + jrp.is_alive.return_value = False + conn = MagicMock() + conn.poll.return_value = True + self._patch_run(task, jrp, conn) + + def test_does_not_mutate_extra_env_with_session_id(self, tmp_path): + extra = {} + task = _task(extra_env=extra) + task.log_stream = MagicMock() + task.log_file = str(tmp_path / 'execution.log') + task.stats_file = str(tmp_path / 'missing.txt') + jrp = MagicMock() + jrp.is_alive.return_value = False + conn = MagicMock() + conn.poll.return_value = True + self._patch_run(task, jrp, conn) + assert extra == {} + assert '__LITHOPS_SESSION_ID' not in extra + assert 'LITHOPS_CONFIG' not in extra + + +class TestGetFunctionAndModules: + + def test_loads_from_storage_without_modules(self): + job = SimpleNamespace( + config=_job_config(), + func_key='func.pickle', + job_key='jk', + ) + payload = pickle.dumps({'func': b'FN', 'module_data': {}}) + storage = MagicMock() + storage.get_func.return_value = payload + assert get_function_and_modules(job, storage) == b'FN' + storage.get_func.assert_called_once_with('func.pickle') + + def test_writes_modules_and_strips_slash(self, tmp_path, monkeypatch): + monkeypatch.setattr('lithops.worker.utils.MODULES_DIR', str(tmp_path)) + job = SimpleNamespace( + config=_job_config(), + func_key='func.pickle', + job_key='jk', + ) + payload = pickle.dumps({ + 'func': b'FN', + 'module_data': { + '/pkg/a.py': bytes_to_b64str(b'aaa'), + '/pkg/b.py': bytes_to_b64str(b'bbb'), + }, + }) + storage = MagicMock() + storage.get_func.return_value = payload + path = os.path.join(str(tmp_path), 'jk') + try: + assert get_function_and_modules(job, storage) == b'FN' + assert (tmp_path / 'jk' / 'pkg' / 'a.py').read_bytes() == b'aaa' + assert (tmp_path / 'jk' / 'pkg' / 'b.py').read_bytes() == b'bbb' + assert path in sys.path + finally: + if path in sys.path: + sys.path.remove(path) + + def test_runtime_include_function_reads_local_file( + self, tmp_path, monkeypatch + ): + monkeypatch.setattr('lithops.worker.utils.SA_INSTALL_DIR', str(tmp_path)) + func_file = tmp_path / 'func.pickle' + func_file.write_bytes(pickle.dumps({'func': b'LOCAL'})) + cfg = _job_config() + cfg['localhost']['runtime_include_function'] = True + job = SimpleNamespace( + config=cfg, func_key='func.pickle', job_key='jk' + ) + assert get_function_and_modules(job, MagicMock()) == b'LOCAL' + + def test_runtime_include_function_uses_posix_install_dir(self): + payload = pickle.dumps({'func': b'LOCAL'}) + cfg = _job_config() + cfg['localhost']['runtime_include_function'] = True + job = SimpleNamespace( + config=cfg, func_key='abc.func.pickle', job_key='jk' + ) + with patch( + 'lithops.worker.utils.SA_INSTALL_DIR', '/opt/lithops' + ), patch('builtins.open', mock_open(read_data=payload)) as opened: + assert get_function_and_modules(job, MagicMock()) == b'LOCAL' + assert opened.call_args[0][0] == '/opt/lithops/abc.func.pickle' + assert '\\' not in opened.call_args[0][0] + + def test_makedirs_uses_exist_ok(self, tmp_path, monkeypatch): + monkeypatch.setattr('lithops.worker.utils.MODULES_DIR', str(tmp_path)) + job = SimpleNamespace( + config=_job_config(), func_key='f', job_key='jk' + ) + payload = pickle.dumps({ + 'func': b'FN', + 'module_data': {'a.py': bytes_to_b64str(b'x')}, + }) + storage = MagicMock() + storage.get_func.return_value = payload + real = os.makedirs + + def _makedirs(path, exist_ok=False): + if not exist_ok: + raise OSError(13, 'denied') + return real(path, exist_ok=exist_ok) + + monkeypatch.setattr(os, 'makedirs', _makedirs) + path = os.path.join(str(tmp_path), 'jk') + try: + assert get_function_and_modules(job, storage) == b'FN' + assert (tmp_path / 'jk' / 'a.py').read_bytes() == b'x' + finally: + if path in sys.path: + sys.path.remove(path) + + +class TestGetFunctionData: + + def test_byte_ranges_slice_aggregated_object(self): + job = SimpleNamespace( + data_key='data.pickle', + data_byte_ranges=[(0, 2), (3, 5)], + ) + storage = MagicMock() + storage.get_data.return_value = b'abcdef' + data = get_function_data(job, storage) + assert data == [b'abc', b'def'] + extra = storage.get_data.call_args.kwargs['extra_get_args'] + assert extra['Range'] == 'bytes=0-5' + + def test_none_byte_ranges_returns_whole_object(self): + job = SimpleNamespace(data_key='data.pickle', data_byte_ranges=None) + storage = MagicMock() + storage.get_data.return_value = b'all' + assert get_function_data(job, storage) == [b'all'] + extra = storage.get_data.call_args.kwargs['extra_get_args'] + assert extra == {} + + def test_empty_byte_ranges_list_returns_whole_object(self): + job = SimpleNamespace(data_key='data.pickle', data_byte_ranges=[]) + storage = MagicMock() + storage.get_data.return_value = b'all' + assert get_function_data(job, storage) == [b'all'] + extra = storage.get_data.call_args.kwargs['extra_get_args'] + assert extra == {} + + def test_payload_data_uses_literal_eval(self): + job = SimpleNamespace(data_key=None, data_byte_strs=["b'abc'", "'x'"]) + assert get_function_data(job, MagicMock()) == [b'abc', 'x'] + + def test_payload_data_bytes_pass_through(self): + job = SimpleNamespace(data_key=None, data_byte_strs=[b'\x80abc']) + assert get_function_data(job, MagicMock()) == [b'\x80abc'] + + def test_payload_data_does_not_eval_expressions(self): + job = SimpleNamespace(data_key=None, data_byte_strs=["1+1"]) + with pytest.raises((ValueError, SyntaxError)): + get_function_data(job, MagicMock()) + + +class TestWorkerUtils: + + def test_custom_redirection_restores_streams(self): + buf = io.StringIO() + old_out, old_err = sys.stdout, sys.stderr + with custom_redirection(buf): + print('hello', end='') + assert sys.stdout is buf + assert sys.stdout is old_out + assert sys.stderr is old_err + assert buf.getvalue() == 'hello' + + def test_log_stream_write_flush_and_valueerror(self): + stream = MagicMock() + ls = LogStream(stream) + ls.write('hi') + stream.write.assert_called_with('hi') + stream.flush.side_effect = ValueError() + ls.flush() + stream.write.side_effect = ValueError() + ls.write('ignored') + assert ls.fileno() == sys.stdout.fileno() + + def test_system_monitor_without_psutil(self, monkeypatch): + monkeypatch.setattr('lithops.worker.utils.psutil_found', False) + mon = SystemMonitor() + mon.start() + mon.stop() + assert mon.get_cpu_info() == {"usage": [], "system": 0, "user": 0} + assert mon.get_network_io() == {"sent": 0, "recv": 0} + assert mon.get_memory_info() == {"rss": 0, "vms": 0, "uss": 0} + + def test_system_monitor_with_psutil(self): + if not psutil_found: + pytest.skip('psutil not installed') + mon = SystemMonitor() + mon.start() + mon.stop() + cpu = mon.get_cpu_info() + assert 'usage' in cpu and 'system' in cpu and 'user' in cpu + net = mon.get_network_io() + assert 'sent' in net and 'recv' in net + mem = mon.get_memory_info() + assert mem['rss'] >= 0 + + def test_get_runtime_metadata(self): + meta = get_runtime_metadata() + assert 'preinstalls' in meta + assert all(len(entry) == 2 for entry in meta['preinstalls']) + assert meta['python_version'] == ( + str(sys.version_info[0]) + "." + str(sys.version_info[1]) + ) + assert 'lithops_version' in meta + + def test_peak_memory_and_disk(self, tmp_path): + mem = peak_memory() + assert mem is None or mem >= 0 + assert free_disk_space(str(tmp_path)) > 0 + + def test_get_memory_usage_non_root_returns_none(self): + if os.geteuid() != 0: + assert get_memory_usage() is None + + def test_memory_monitor_sends_peak_on_poll(self): + conn = MagicMock() + conn.poll.return_value = True + with patch( + 'lithops.worker.utils.get_memory_usage', return_value=10 + ): + memory_monitor_worker(conn, delay=0) + conn.send.assert_called_once() + + def test_memory_monitor_breaks_when_usage_is_none(self): + conn = MagicMock() + conn.poll.side_effect = [False, True] + with patch( + 'lithops.worker.utils.get_memory_usage', return_value=None + ): + memory_monitor_worker(conn, delay=0) + conn.send.assert_called_once_with(0) + + +class TestCallStatus: + + def test_create_call_status_storage_and_rabbitmq(self): + job = _task() + st = create_call_status(job, MagicMock()) + assert isinstance(st, StorageCallStatus) + job.config['lithops']['monitoring'] = 'rabbitmq' + rb = create_call_status(job, MagicMock()) + assert isinstance(rb, RabbitmqCallStatus) + + def test_warm_container_flag(self, monkeypatch): + monkeypatch.delenv('WARM_CONTAINER', raising=False) + job = _task() + first = CallStatus(job, MagicMock()) + assert first.status['worker_cold_start'] is True + assert os.environ['WARM_CONTAINER'] == 'True' + second = CallStatus(job, MagicMock()) + assert second.status['worker_cold_start'] is False + + def test_warm_container_invalid_value_is_cold_start(self, monkeypatch): + monkeypatch.setenv('WARM_CONTAINER', 'maybe') + status = CallStatus(_task(), MagicMock()) + assert status.status['worker_cold_start'] is True + assert os.environ['WARM_CONTAINER'] == 'True' + + def test_storage_init_and_end_events(self): + import json + storage = MagicMock() + status = StorageCallStatus(_task(), storage) + status.send_init_event() + assert storage.put_data.call_args[0][1] == '' + status.add('foo', 1) + status.send_finish_event() + body = storage.put_data.call_args[0][1] + payload = json.loads(body) + assert payload['type'] == '__end__' + assert payload['foo'] == 1 + + def test_rabbitmq_end_also_writes_storage(self): + job = _task() + job.executor_id = 'a-b-c-d' + # One publish per queue in the chain the client sent + job.monitoring_queues = ['lithops-a-b', 'lithops-a-b-c-d'] + job.config = _job_config(monitoring='rabbitmq') + storage = MagicMock() + channel = MagicMock() + conn = MagicMock() + conn.channel.return_value = channel + status = RabbitmqCallStatus(job, storage) + status.status['type'] = '__end__' + status.status['activation_id'] = 'act' + with patch( + 'lithops.worker.status.pika.BlockingConnection', return_value=conn + ): + status._send() + assert channel.basic_publish.call_count == 2 + assert storage.put_data.called + + def test_rabbitmq_publishes_to_the_queues_it_was_given(self): + job = _task() + job.executor_id = 'sess-0-M000-00000-0' + job.monitoring_queues = [ + 'lithops-sess-0', 'lithops-sess-0-M000-00000-0' + ] + job.config = _job_config(monitoring='rabbitmq') + status = RabbitmqCallStatus(job, MagicMock()) + assert status._queue_names() == job.monitoring_queues + + def test_rabbitmq_falls_back_to_its_own_queue(self): + # A payload with no chain can only reach this executor's own queue + job = _task() + job.executor_id = 'sess-0-M000-00000-0' + job.config = _job_config(monitoring='rabbitmq') + status = RabbitmqCallStatus(job, MagicMock()) + assert status._queue_names() == ['lithops-sess-0-M000-00000-0'] + + def test_rabbitmq_gives_up_after_five_failures(self): + job = _task() + job.executor_id = 'a-b' + job.config = _job_config(monitoring='rabbitmq') + storage = MagicMock() + status = RabbitmqCallStatus(job, storage) + status.status['type'] = '__init__' + with patch( + 'lithops.worker.status.pika.BlockingConnection', + side_effect=Exception('down'), + ): + test_thread = threading.current_thread() + sleeps = [] + + def sleep(_seconds): + if threading.current_thread() is test_thread: + sleeps.append(_seconds) + + with patch('lithops.worker.status.time.sleep', side_effect=sleep): + status._send() + assert len(sleeps) == 5 + storage.put_data.assert_not_called() + + +class TestJobStatsAndPrepareArgs: + + def test_job_stats_write(self, tmp_path): + path = tmp_path / 'stats.txt' + stats = JobStats(str(path)) + stats.write('k', 1.5) + stats.write('s', 'x') + stats.__del__() + assert path.read_text() == 'k 1.5\ns x\n' + + def test_prepare_args_kwargs_only(self): + def f(a, b=1): + return a + b + args, kwargs = _prepare_args(f, {'a': 2, 'b': 3}) + assert args == () + assert kwargs == {'a': 2, 'b': 3} + assert f(*args, **kwargs) == 5 + + def test_prepare_args_varargs_empty_list_is_kept(self): + def f(*args, **kwargs): + return args, kwargs + args, kwargs = _prepare_args(f, {'args': [], 'kwargs': {}, 'x': 1}) + assert args == [] + assert kwargs == {'x': 1} + + def test_prepare_args_custom_var_names(self): + def f(*xs, **kw): + return xs, kw + args, kwargs = _prepare_args( + f, {'xs': (1, 2), 'kw': {'a': 3}, 'b': 4} + ) + assert args == (1, 2) + assert kwargs == {'a': 3, 'b': 4} + + +class TestJobRunner: + + @pytest.fixture(autouse=True) + def _session_id(self, monkeypatch, tmp_path): + monkeypatch.setenv('__LITHOPS_SESSION_ID', 'sid-1') + self.stats = str(tmp_path / 'stats.txt') + + def _runner(self, func, data, **job_kwargs): + job = _task( + func=pickle.dumps(func), + data=pickle.dumps(data), + stats_file=self.stats, + config=_job_config(telemetry=False), + **job_kwargs + ) + conn = MagicMock() + storage = MagicMock() + storage.backend = 'localhost' + storage.storage = MagicMock() + return JobRunner(job, conn, storage) + + def test_run_small_result_stored_in_stats(self): + jr = self._runner(_add, {'x': 1}) + jr.run() + jr.jobrunner_conn.send.assert_called_with('Finished') + text = open(self.stats).read() + assert 'func_result_size' in text + assert 'result' in text + jr.internal_storage.put_data.assert_not_called() + + def test_run_large_result_uploaded(self): + jr = self._runner(_big, {'x': 1}) + jr.run() + jr.internal_storage.put_data.assert_called_once() + text = open(self.stats).read() + assert 'worker_result_upload_time' in text + + def test_run_none_result_skips_upload(self): + jr = self._runner(_none, {'x': 1}) + jr.run() + jr.internal_storage.put_data.assert_not_called() + + def test_run_exception_records_exc_info(self): + jr = self._runner(_boom, {'x': 1}) + jr.run() + text = open(self.stats).read() + assert 'exception True' in text + assert 'exc_info' in text + jr.jobrunner_conn.send.assert_called_with('Finished') + + def test_fill_optional_args_id_and_storage(self): + jr = self._runner(_echo, {'x': 1}, call_id='00007') + data = {'x': 1} + jr._fill_optional_args(_with_id_and_storage, data) + assert data['id'] == 7 + assert data['storage'] is jr.internal_storage.storage + + def test_fill_optional_args_missing_ibm_cos_and_rabbitmq(self): + def f(ibm_cos): + pass + + def g(rabbitmq): + pass + jr = self._runner(_echo, {'x': 1}) + jr.lithops_config.pop('rabbitmq', None) + with pytest.raises(Exception, match='ibm_cos'): + jr._fill_optional_args(f, {}) + with pytest.raises(Exception, match='rabbitmq'): + jr._fill_optional_args(g, {}) + + def test_fill_optional_args_ibm_cos_same_backend(self): + def f(ibm_cos): + pass + jr = self._runner(_echo, {'x': 1}) + jr.lithops_config['ibm_cos'] = {} + jr.internal_storage.backend = 'ibm_cos' + jr.internal_storage.get_client.return_value = 'client' + data = {} + jr._fill_optional_args(f, data) + assert data['ibm_cos'] == 'client' + + def test_fill_optional_args_ibm_cos_other_backend(self): + def f(ibm_cos): + pass + jr = self._runner(_echo, {'x': 1}) + jr.lithops_config['ibm_cos'] = {} + data = {} + with patch('lithops.worker.jobrunner.Storage') as st: + st.return_value.get_client.return_value = 'other' + jr._fill_optional_args(f, data) + assert data['ibm_cos'] == 'other' + + def test_fill_optional_args_future_chaining(self): + jr = self._runner(_echo, {'x': 1}) + future = MagicMock() + future.result.return_value = 9 + data = {'future': future} + jr._fill_optional_args(_echo, data) + assert data['x'] == 9 + assert 'future' not in data + + def test_wait_futures_replaces_first_value(self): + jr = self._runner(_echo, {'x': 1}) + done = MagicMock(done=True, futures=False) + done.result.return_value = 5 + skip = MagicMock(done=True, futures=True) + data = {'results': [done, skip]} + with patch('lithops.worker.jobrunner.wait'): + jr._wait_futures(data) + assert data['results'] == [5] + + def test_load_object_from_path(self, tmp_path): + path = tmp_path / 'obj.bin' + path.write_bytes(b'abcdefghij') + obj = CloudObjectLocal(str(path)) + obj.data_byte_range = (2, 6) + obj.chunk_size = 5 + obj.newline = None + obj.part = 1 + obj.total_parts = 1 + jr = self._runner(_obj_fn, {'obj': obj}) + jr._load_object({'obj': obj}) + assert obj.data_stream.read() == b'cdefg' + assert obj.data_byte_range == (2, 6) + + def test_load_object_without_range_sets_full_chunk(self, tmp_path): + path = tmp_path / 'obj.bin' + path.write_bytes(b'abc') + obj = CloudObjectLocal(str(path)) + obj.data_byte_range = None + obj.chunk_size = 10 + obj.part = 1 + obj.total_parts = 1 + jr = self._runner(_obj_fn, {'obj': obj}) + jr._load_object({'obj': obj}) + assert obj.data_byte_range == (0, 9) + assert obj.data_stream.read() == b'abc' + + def test_load_object_url_and_storage(self): + obj = CloudObjectUrl('http://example.com/a') + obj.data_byte_range = (0, 10) + obj.chunk_size = 5 + obj.newline = '\n' + obj.part = 1 + obj.total_parts = 2 + raw = io.BytesIO(b'hello') + resp = MagicMock(raw=raw) + jr = self._runner(_obj_fn, {'obj': obj}) + with patch( + 'lithops.worker.jobrunner.requests.get', return_value=resp + ): + jr._load_object({'obj': obj}) + assert obj.data_byte_range == (0, 4) + + def test_load_object_cloudobject_other_backend(self): + obj = CloudObject('aws_s3', 'b', 'k') + obj.data_byte_range = None + obj.chunk_size = 3 + obj.part = 1 + obj.total_parts = 1 + jr = self._runner(_obj_fn, {'obj': obj}) + with patch('lithops.worker.jobrunner.Storage') as st: + st.return_value.get_object.return_value = io.BytesIO(b'xyz') + jr._load_object({'obj': obj}) + st.assert_called_once() + + def test_new_futures_list_skips_result_upload(self): + jr = self._runner(_return_futures_list, {'x': 1}) + jr.run() + text = open(self.stats).read() + assert 'new_futures' in text + jr.internal_storage.put_data.assert_not_called() + + def test_reduce_job_waits_for_futures(self, monkeypatch): + monkeypatch.setenv('__LITHOPS_REDUCE_JOB', 'True') + jr = self._runner(_reduce_fn, {'results': []}) + with patch.object(jr, '_wait_futures') as wait_f: + jr.run() + wait_f.assert_called_once() + + def test_object_processing_loads_object(self, tmp_path): + path = tmp_path / 'o.bin' + path.write_bytes(b'abcd') + obj = CloudObjectLocal(str(path)) + obj.data_byte_range = None + obj.chunk_size = 4 + obj.part = 1 + obj.total_parts = 1 + jr = self._runner(_obj_fn, {'obj': obj}) + with patch.object(jr, '_load_object') as load: + jr.run() + load.assert_called_once() + + def test_prepost_hooks(self, monkeypatch): + calls = [] + + def pre(): + calls.append('pre') + + def post(): + calls.append('post') + + monkeypatch.setenv('PRE_RUN', 'pre') + monkeypatch.setenv('POST_RUN', 'post') + jr = self._runner(_echo, {'x': 1}) + with patch( + 'lithops.worker.jobrunner.locate', side_effect=[pre, post] + ): + jr.run() + assert calls == ['pre', 'post'] + + def test_callable_class_function_name(self): + jr = self._runner(_Adder(), {'x': 1}) + jr.run() + text = open(self.stats).read() + assert 'func_result_size' in text + + +class TestFunctionInvoker: + + def test_function_invoker_wires_handlers(self, monkeypatch): + monkeypatch.delenv('LITHOPS_WORKER', raising=False) + payload = { + 'config': _job_config( + monitoring='storage', backend='aws_lambda' + ), + 'job': { + 'job_key': 'jk', + 'executor_id': 'ex', + 'job_id': 'j0', + 'chunksize': 1, + }, + } + payload['config']['aws_lambda'] = {} + invoker = MagicMock() + with patch( + 'lithops.worker.invoker.extract_storage_config', return_value={} + ): + with patch('lithops.worker.invoker.InternalStorage'): + with patch( + 'lithops.worker.invoker.extract_serverless_config', + return_value={}, + ): + with patch('lithops.worker.invoker.ServerlessHandler'): + with patch('lithops.worker.invoker.JobMonitor'): + with patch( + 'lithops.worker.invoker.FaaSRemoteInvoker', + return_value=invoker, + ): + function_invoker(payload) + invoker.run_job.assert_called_once() + assert os.environ['LITHOPS_WORKER'] == 'True' + assert payload['config']['aws_lambda']['invoke_pool_threads'] == 128 + + def test_remote_invoker_run_job_drains_then_stops_waiting(self): + from lithops.worker.invoker import FaaSRemoteInvoker + inv = FaaSRemoteInvoker.__new__(FaaSRemoteInvoker) + inv.job_monitor = MagicMock() + inv.pending_calls_q = MagicMock() + inv.pending_calls_q.qsize.side_effect = [2, 0] + inv.stop = MagicMock() + inv._run_job = MagicMock(return_value=['f']) + job = SimpleNamespace(job_id='j0', chunksize=1) + test_thread = threading.current_thread() + sleeps = [] + + def sleep(_seconds): + if threading.current_thread() is test_thread: + sleeps.append(_seconds) + + with patch('lithops.worker.invoker.time.sleep', side_effect=sleep): + inv.run_job(job) + inv.job_monitor.start.assert_called_once() + inv.job_monitor.stop.assert_called_once() + # Waits for the invocations in flight instead of sleeping on a guess + inv.stop.assert_called_once_with(wait=True) + assert sleeps == [1] diff --git a/lithops/util/ibm_token_manager.py b/lithops/util/ibm_token_manager.py index a9c11dcdb..0d2d54af7 100644 --- a/lithops/util/ibm_token_manager.py +++ b/lithops/util/ibm_token_manager.py @@ -15,120 +15,133 @@ # import os +import time import logging from datetime import datetime, timezone +from typing import Optional, Tuple + from ibm_botocore.credentials import DefaultTokenManager +from ibm_cloud_sdk_core.authenticators import IAMAuthenticator from lithops.config import load_yaml_config, dump_yaml_config from lithops.constants import CACHE_DIR -from ibm_cloud_sdk_core.authenticators import IAMAuthenticator logger = logging.getLogger(__name__) - -# The token will be considered expired 20 minutes before its actual expiration time +# The token will be considered expired 20 minutes before its expiration time EXPIRY_MINUTES = 20 +# How long an IBM Cloud token lasts, used only when the library does not tell +# us the real expiry time +DEFAULT_TOKEN_LIFETIME = 60 * 60 -class IBMTokenManager: - - TOEKN_FILE = None - TYPE = None - def __init__(self, ibm_api_key, token=None, token_expiry_time=None): +class IBMTokenManager: + """ + Keeps an IBM Cloud token valid, caching it in a local file so that it can + be reused across executions. Subclasses provide the token generation + """ + + TOKEN_FILE: Optional[str] = None + TYPE: Optional[str] = None + + def __init__( + self, + ibm_api_key: str, + token: Optional[str] = None, + token_expiry_time: Optional[int] = None, + ): self.ibm_api_key = ibm_api_key self.token = token self.expiry_time = token_expiry_time + token_source = 'the configuration' - if not self.token and os.path.exists(self.TOEKN_FILE): - token_data = load_yaml_config(self.TOEKN_FILE) + if not self.token and self.TOKEN_FILE and os.path.exists(self.TOKEN_FILE): + token_data = load_yaml_config(self.TOKEN_FILE) self.token = token_data.get('token') self.expiry_time = token_data.get('expiry_time') + token_source = 'local cache' if not self._is_token_expired(): - logger.debug(f"Reusing {self.TYPE} token from local cache") + logger.debug(f"Reusing {self.TYPE} token from {token_source}") self._log_remaining_time() - def _is_token_expired(self): + def _is_token_expired(self) -> bool: """ - Checks if a token already expired + Checks whether the token is missing, expired, or about to expire """ return self._get_token_minutes_left() < EXPIRY_MINUTES - def _get_token_minutes_left(self): - """ - Gets the remaining minutes in which the current token is valid - """ + def _get_token_minutes_left(self) -> int: + """Gets the minutes the current token is still valid for""" if not self.expiry_time: return 0 expiry_time = datetime.fromtimestamp(self.expiry_time, tz=timezone.utc) - return max(0, int((expiry_time - datetime.now(timezone.utc)).total_seconds() / 60.0)) + remaining = (expiry_time - datetime.now(timezone.utc)).total_seconds() + return max(0, int(remaining / 60.0)) - def _generate_new_token(self): - """ - Generates a new token - """ + def _generate_new_token(self) -> None: + """Requests a new token and stores it with its expiry time""" raise NotImplementedError() - def _log_remaining_time(self): - """ - Logs the remaining time of the token - """ - minutes_left = self._get_token_minutes_left() - expiry_time = datetime.fromtimestamp(self.expiry_time) - logger.debug(f"{self.TYPE} token expiry time: {expiry_time} - Minutes left: {minutes_left}") - - def _dump_token_data(self): - """ - Dumps the token into a local cache file - """ + def _log_remaining_time(self) -> None: + expiry_time = datetime.fromtimestamp(self.expiry_time, tz=timezone.utc) + logger.debug( + f"{self.TYPE} token expiry time: {expiry_time} - " + f"Minutes left: {self._get_token_minutes_left()}" + ) + + def _dump_token_data(self) -> None: + if not self.TOKEN_FILE: + return token_data = {'token': self.token, 'expiry_time': self.expiry_time} - dump_yaml_config(self.TOEKN_FILE, token_data) + dump_yaml_config(self.TOKEN_FILE, token_data) - def refresh_token(self): - """ - Forces to create a new token - """ + def refresh_token(self) -> Tuple[Optional[str], Optional[int]]: + """Generates a new token, caches it, and returns it""" self._generate_new_token() self._dump_token_data() self._log_remaining_time() - return self.token, self.expiry_time - def get_token(self): - """ - Gets the current token or creates a new one if expired - """ + def get_token(self) -> Tuple[Optional[str], Optional[int]]: + """Gets the current token, refreshing it first if it is expired""" if self._is_token_expired(): self.refresh_token() - return self.token, self.expiry_time class COSTokenManager(IBMTokenManager): + """Token manager for IBM Cloud Object Storage""" - TOEKN_FILE = os.path.join(CACHE_DIR, 'ibm_cos', 'token') + TOKEN_FILE = os.path.join(CACHE_DIR, 'ibm_cos', 'token') TYPE = 'COS' - def _generate_new_token(self): - """ - Generates a new COS token - """ + def _generate_new_token(self) -> None: logger.debug("Requesting new COS token") token_manager = DefaultTokenManager(api_key_id=self.ibm_api_key) self.token = token_manager.get_token() - self.expiry_time = int(token_manager._expiry_time.timestamp()) + # ibm_botocore exposes the expiry time only as a private attribute, so + # a library upgrade can take it away + expiry_time = getattr(token_manager, '_expiry_time', None) + if expiry_time is None: + logger.warning( + "ibm_botocore no longer reports the token expiry time, " + f"assuming the standard {DEFAULT_TOKEN_LIFETIME // 60} " + "minute lifetime" + ) + self.expiry_time = int(time.time()) + DEFAULT_TOKEN_LIFETIME + else: + self.expiry_time = int(expiry_time.timestamp()) class IAMTokenManager(IBMTokenManager): + """Token manager for IBM Cloud IAM""" - TOEKN_FILE = os.path.join(CACHE_DIR, 'ibm_iam', 'token') + TOKEN_FILE = os.path.join(CACHE_DIR, 'ibm_iam', 'token') TYPE = 'IAM' - def _generate_new_token(self): - """ - Generates a new IAM token - """ + def _generate_new_token(self) -> None: logger.debug("Requesting new IAM token") auth = IAMAuthenticator(self.ibm_api_key) self.token = auth.token_manager.get_token() diff --git a/lithops/util/joblib/__init__.py b/lithops/util/joblib/__init__.py index 09cd9bed4..7af838f9b 100644 --- a/lithops/util/joblib/__init__.py +++ b/lithops/util/joblib/__init__.py @@ -2,13 +2,9 @@ def register_lithops(): - """ Register Lithops Backend to be called with parallel_backend("lithops"). """ - try: - from lithops.util.joblib.lithops_backend import LithopsBackend - register_parallel_backend("lithops", LithopsBackend) - except ImportError: - msg = ("To use the Lithops backend you must install lithops.") - raise ImportError(msg) + """Register Lithops Backend to be called with parallel_backend("lithops").""" + from lithops.util.joblib.lithops_backend import LithopsBackend + register_parallel_backend("lithops", LithopsBackend) __all__ = ["register_lithops"] diff --git a/lithops/util/joblib/lithops_backend.py b/lithops/util/joblib/lithops_backend.py index c742f80c8..497195b5a 100644 --- a/lithops/util/joblib/lithops_backend.py +++ b/lithops/util/joblib/lithops_backend.py @@ -17,34 +17,41 @@ import logging import os import pickle +import threading +from concurrent.futures import Future, ThreadPoolExecutor +from multiprocessing.pool import ThreadPool +from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple + import diskcache from numpy import ndarray -from multiprocessing.pool import ThreadPool -from concurrent.futures import ThreadPoolExecutor -from typing import Any, Dict, Optional from joblib import Parallel from joblib._parallel_backends import MultiprocessingBackend from joblib.pool import PicklingPool -from joblib.parallel import register_parallel_backend from lithops.multiprocessing import Pool, cpu_count +from lithops.multiprocessing import config as mp_config from lithops.constants import LITHOPS_TEMP_DIR from lithops.storage import Storage logger = logging.getLogger(__name__) +# A call, as joblib hands it over: (function, args, kwargs), optionally +# followed by the positions of the arguments replaced by a cloud object +Call = Tuple -def register_lithops(): - """ Register Lithops Backend to be called with parallel_backend("lithops"). """ - register_parallel_backend("lithops", LithopsBackend) +# Upper bound for the threads that upload and download the shared arguments: +# there is one argument per thread, and a batch can hold a great many +MAX_UPLOAD_THREADS = 32 + +# Tells a cached None apart from a key that is not cached +_CACHE_MISS = object() class LithopsBackend(MultiprocessingBackend): - """A ParallelBackend which will use a multiprocessing.Pool. - Will introduce some communication and memory overhead when exchanging - input and output data with the with the worker Python processes. - However, does not suffer from the Python Global Interpreter Lock. + """ + joblib backend that runs the tasks of a batch through Lithops instead of + a local process pool, uploading the arguments they share only once """ supports_timeout = True @@ -61,6 +68,11 @@ def __init__( self.lithops_args = lithops_args self.eff_n_jobs = None self.prefer = None + + if lithops_args: + # The batches run on lithops.multiprocessing, which takes its + # executor arguments from this process-wide parameter + mp_config.set_parameter(mp_config.LITHOPS_CONFIG, lithops_args) super().__init__( nesting_level=nesting_level, inner_max_num_threads=inner_max_num_threads, @@ -75,158 +87,237 @@ def configure( require: Optional[str] = None, **memmappingpool_args ): - """Make Lithops Pool the father class of PicklingPool. PicklingPool is a - father class that inherits Pool from multiprocessing.pool. The next - line is a patch, which changes the inheritance of Pool to be from - lithops.multiprocessing.pool + """ + Configures the backend, making the Lithops Pool the one that joblib + instantiates to run the tasks """ self.prefer = prefer + # PicklingPool inherits Pool from multiprocessing.pool. This patch + # changes that inheritance to lithops.multiprocessing.Pool PicklingPool.__bases__ = (Pool,) if n_jobs == -1: n_jobs = self.effective_n_jobs(n_jobs) - eff_n_jobs = super(LithopsBackend, self).configure( + return super().configure( n_jobs, parallel, prefer, require, **memmappingpool_args ) - return eff_n_jobs - def effective_n_jobs(self, n_jobs): - eff_n_jobs = super(LithopsBackend, self).effective_n_jobs(n_jobs) + def effective_n_jobs(self, n_jobs: int) -> int: + """ + Resolves how many jobs to run in parallel, asking the Lithops backend + only once for the CPUs that -1 stands for + """ + eff_n_jobs = super().effective_n_jobs(n_jobs) if n_jobs == -1: self.eff_n_jobs = self.eff_n_jobs or cpu_count() eff_n_jobs = self.eff_n_jobs return eff_n_jobs def start_call(self): - """This is a workaround to make "batch size" working properly - and invoke all the tasks using a single map() instead of - individual apply_async()""" + """Forces a single map() instead of one apply_async() per task""" self.parallel._cached_effective_n_jobs = 1 self.parallel.pre_dispatch = 'all' - def compute_batch_size(self): + def compute_batch_size(self) -> int: + """Keeps every task in one batch, so that a call is a single map""" return int(1e6) - def apply_async(self, func, callback=None): - """Schedule a func to be run""" + def submit(self, func, callback=None): + """ + Schedules a batch of calls, uploading the arguments they have in + common only once. + + joblib renamed this hook from apply_async to submit, and the + multiprocessing backend this one extends carries its own submit, so + without this override joblib would hand the batch straight to the pool + wrapped in a class the pool does not know what to do with + """ mem_opt_calls = find_shared_objects(func.items) + pool = self._get_pool() if self.prefer == "threads": - return self._get_pool().apply_async(handle_call_threads, (mem_opt_calls, ), callback=callback) - else: - return self._get_pool().starmap_async(handle_call_process, mem_opt_calls, callback=callback) + return pool.apply_async( + handle_call_threads, (mem_opt_calls,), callback=callback + ) + return pool.starmap_async( + handle_call_process, mem_opt_calls, callback=callback + ) + # The name joblib used before 1.4 + apply_async = submit + + +def _wait_all(futures: Iterable[Future]) -> None: + """Waits for every future, re-raising whatever they raised""" + for fut in futures: + fut.result() -def find_shared_objects(calls): - # find and annotate repeated arguments - logger.info('Optimizing shared data between tasks') +def _storage_for_the_pool() -> Storage: + """ + Builds a storage client from the parameters the pool runs with, so that a + shared argument lands where the workers of that pool will look for it and + not in whatever storage this machine has configured by default + """ + lithops_conf = mp_config.get_parameter(mp_config.LITHOPS_CONFIG) or {} + return Storage( + config=lithops_conf.get('config'), + backend=lithops_conf.get('storage'), + ) + + +def _index_arguments_by_identity(calls: Sequence[Call]) -> Dict[int, List]: + """ + Groups every argument of every call by object identity. Each entry holds + the object itself followed by the (call, position) pairs it appears in, + where a position is an index for an arg and a name for a kwarg + """ record = {} for i, call in enumerate(calls): - for j, arg in enumerate(call[1]): - if id(arg) in record: - record[id(arg)].append((i, j)) - else: - record[id(arg)] = [arg, (i, j)] - - for k, v in call[2].items(): - if id(v) in record: - record[id(v)].append((i, k)) - else: - record[id(v)] = [v, (i, k)] - - # If we found multiple occurrences of one object, then - # store it in shared memory, pass a proxy as a value - calls = [list(item) for item in calls] + arguments = list(enumerate(call[1])) + list(call[2].items()) + for idx_or_key, arg in arguments: + record.setdefault(id(arg), [arg]).append((i, idx_or_key)) + return record + + +def _proxy_argument(call: List, idx_or_key, cloud_object) -> None: + """ + Replaces one argument of a call with a cloud object, and records its + position so that the worker knows which arguments to fetch back + """ + if isinstance(idx_or_key, str): + call[2][idx_or_key] = cloud_object + else: + args_as_list = list(call[1]) + args_as_list[idx_or_key] = cloud_object + call[1] = tuple(args_as_list) + + # The 4th element only exists once a first argument has been proxied + try: + call[3].append(idx_or_key) + except IndexError: + call.append([idx_or_key]) + + +def find_shared_objects(calls: Sequence[Call]) -> List[Call]: + """ + Replaces the arguments that several calls share with a proxy to a single + cloud object, so that they travel to the workers only once + """ + logger.info('Optimizing shared data between tasks') - storage = Storage() - thread_pool = ThreadPoolExecutor(max_workers=len(record)) + record = _index_arguments_by_identity(calls) + calls = [list(item) for item in calls] + if not record: + return [tuple(item) for item in calls] + + storage = None + storage_lock = threading.Lock() + # Two shared arguments of the same call are proxied by two threads, and + # each one rewrites the args tuple of every call it appears in + calls_lock = threading.Lock() + + def get_storage(): + # Created on first use and shared, so that the uploading threads do + # not build one client each + nonlocal storage + with storage_lock: + if storage is None: + storage = _storage_for_the_pool() + return storage def put_arg_obj(positions): obj = positions.pop(0) - if len(positions) > 1 and consider_sharing(obj): - logger.debug('Proxying {}'.format(type(obj))) - obj_bin = pickle.dumps(obj) - cloud_object = storage.put_cloudobject(obj_bin) - - for pos in positions: - call_n, idx_or_key = pos - call = calls[call_n] - - if isinstance(idx_or_key, str): - call[2][idx_or_key] = cloud_object - else: - args_as_list = list(call[1]) - args_as_list[idx_or_key] = cloud_object - call[1] = tuple(args_as_list) - - try: - call[3].append(idx_or_key) - except IndexError: - call.append([idx_or_key]) - - fut = [] - for positions in record.values(): - f = thread_pool.submit(put_arg_obj, positions) - fut.append(f) - [f.result() for f in fut] + if len(positions) <= 1 or not consider_sharing(obj): + return - return [tuple(item) for item in calls] + logger.debug(f'Proxying {type(obj)}') + obj_bin = pickle.dumps(obj) + cloud_object = get_storage().put_cloudobject(obj_bin) + with calls_lock: + for call_n, idx_or_key in positions: + _proxy_argument(calls[call_n], idx_or_key, cloud_object) -def handle_call_threads(mem_opt_calls): - with ThreadPool(processes=len(mem_opt_calls)) as pool: - results = pool.starmap(handle_call_process, mem_opt_calls) + workers = min(len(record), MAX_UPLOAD_THREADS) + with ThreadPoolExecutor(max_workers=workers) as thread_pool: + _wait_all([ + thread_pool.submit(put_arg_obj, positions) + for positions in record.values() + ]) - return list(results) + return [tuple(item) for item in calls] -def handle_call_process(func, args, kwargs, proxy_positions=[]): - if len(proxy_positions) > 0: - args, kwargs = replace_with_values(args, kwargs, proxy_positions) +def handle_call_threads(mem_opt_calls: Sequence[Call]) -> List[Any]: + """Runs a whole batch of calls in this worker, one thread each""" + with ThreadPool(processes=max(1, len(mem_opt_calls))) as pool: + return list(pool.starmap(handle_call_process, mem_opt_calls)) + +def handle_call_process( + func: Callable, + args: Tuple, + kwargs: Dict[str, Any], + proxy_positions: Optional[List] = None +) -> Any: + """Runs a single call, fetching the arguments that were proxied""" + if proxy_positions: + args, kwargs = replace_with_values(args, kwargs, proxy_positions) return func(*args, **kwargs) -def replace_with_values(args, kwargs, proxy_positions): +def replace_with_values( + args: Tuple, + kwargs: Dict[str, Any], + proxy_positions: List +) -> Tuple[List, Dict[str, Any]]: + """ + Downloads the cloud objects standing in for the proxied arguments, using + a local disk cache shared by every task that runs in the same worker + """ args_as_list = list(args) - thread_pool = ThreadPoolExecutor(max_workers=len(proxy_positions)) - cache = diskcache.Cache(os.path.join(LITHOPS_TEMP_DIR, 'cache')) + cache_dir = os.path.join(LITHOPS_TEMP_DIR, 'cache') - def get_arg_obj(idx_or_key): + def get_arg_obj(idx_or_key, cache): if isinstance(idx_or_key, str): obj_id = kwargs[idx_or_key] else: obj_id = args_as_list[idx_or_key] - if obj_id in cache: - logger.debug('Get {} (arg {}) from cache'.format(obj_id, idx_or_key)) - obj = cache[obj_id] - else: - logger.debug('Get {} (arg {}) from storage'.format(obj_id, idx_or_key)) + # Read in one call: asking whether the key is there and then reading + # it is a race. Every task of this runtime shares the cache directory, + # and a value too big to sit inline is a file of its own, so the row + # can be there while the file is not readable yet + obj = cache.get(obj_id, default=_CACHE_MISS) + if obj is _CACHE_MISS: + logger.debug(f'Get {obj_id} (arg {idx_or_key}) from storage') storage = Storage() obj_bin = storage.get_cloudobject(obj_id) obj = pickle.loads(obj_bin) cache[obj_id] = obj + else: + logger.debug(f'Get {obj_id} (arg {idx_or_key}) from cache') if isinstance(idx_or_key, str): kwargs[idx_or_key] = obj else: args_as_list[idx_or_key] = obj - fut = [] - for idx_or_key in proxy_positions: - f = thread_pool.submit(get_arg_obj, idx_or_key) - fut.append(f) - [f.result() for f in fut] + with diskcache.Cache(cache_dir) as cache: + workers = min(max(1, len(proxy_positions)), MAX_UPLOAD_THREADS) + with ThreadPoolExecutor(max_workers=workers) as thread_pool: + _wait_all([ + thread_pool.submit(get_arg_obj, idx_or_key, cache) + for idx_or_key in proxy_positions + ]) return args_as_list, kwargs -def consider_sharing(obj): - if isinstance(obj, (ndarray, list)): # TODO: some heuristic - return True - return False +def consider_sharing(obj: Any) -> bool: + """Tells whether an object is worth uploading as a shared cloud object""" + return isinstance(obj, (ndarray, list)) diff --git a/lithops/util/metrics.py b/lithops/util/metrics.py index 035ec51bc..d3995c259 100644 --- a/lithops/util/metrics.py +++ b/lithops/util/metrics.py @@ -1,31 +1,48 @@ -import requests import logging import os +from typing import Any, Dict, Optional, Sequence, Tuple + +import requests logger = logging.getLogger(__name__) +_DEFAULT_INSTANCE = 'lithops' -class PrometheusExporter(): - def __init__(self, enabled, config): - """ Prometheus exporter for sending metrics to an API Gateway""" +class PrometheusExporter: + """ + Pushes Lithops metrics to a Prometheus pushgateway sitting behind an API + Gateway. Does nothing unless it is enabled and a gateway is configured + """ + + def __init__(self, enabled: bool, config: Optional[Dict[str, Any]]): self.enabled = enabled self.apigateway = config.get('apigateway') if config else None - self.job = 'lithops' - self.instance = os.environ['__LITHOPS_SESSION_ID'].split('-')[0] - - def send_metric(self, name, value, type, labels): - """Send a metric to prometheus""" - - if self.enabled and self.apigateway: - dim = 'job/{}/instance/{}'.format(self.job, self.instance) - for key, val in labels: - dim += '/%s/%s' % (key, val) - url = '/'.join([self.apigateway, 'metrics', dim]) - logger.debug('Sending metric "{} {} ({})" to {}'.format(name, value, type, url)) - - try: - requests.post(url, data='# TYPE %s %s\n%s %s\n' % (name, type, name, value)) - except Exception as e: - logger.error(e) + session_id = os.environ.get('__LITHOPS_SESSION_ID', _DEFAULT_INSTANCE) + self.instance = session_id.split('-')[0] + + def send_metric( + self, + name: str, + value: Any, + type: str, + labels: Sequence[Tuple[str, Any]], + ) -> None: + """ + Sends a single metric, with the labels appended to the pushgateway + grouping key. Errors are logged and swallowed: metrics are optional + """ + if not (self.enabled and self.apigateway): + return + + dim = f'job/{self.job}/instance/{self.instance}' + for key, val in labels: + dim += f'/{key}/{val}' + url = '/'.join([self.apigateway, 'metrics', dim]) + logger.debug(f'Sending metric "{name} {value} ({type})" to {url}') + + try: + requests.post(url, data=f'# TYPE {name} {type}\n{name} {value}\n') + except Exception as exc: + logger.error(exc) diff --git a/lithops/util/ssh_client.py b/lithops/util/ssh_client.py index 92c35fd07..9231fe3df 100644 --- a/lithops/util/ssh_client.py +++ b/lithops/util/ssh_client.py @@ -1,6 +1,9 @@ -import paramiko import logging import os +from contextlib import contextmanager +from typing import Any, Dict, List, Optional, Tuple + +import paramiko logger = logging.getLogger(__name__) @@ -9,10 +12,13 @@ for _log_name in ('paramiko', 'paramiko.transport', 'paramiko.client'): logging.getLogger(_log_name).setLevel(logging.CRITICAL) +_DEFAULT_KEY = os.path.expanduser('~/.ssh/id_rsa') -def ssh_boot_status_message(err): + +def ssh_boot_status_message(err: BaseException) -> str: """ - Map transient SSH errors during VM boot to a short user-facing status. + Maps a transient SSH error raised while a VM boots to a short status + message, falling back to the error itself when it is not a known one """ msg = str(err).lower() if 'timed out' in msg or 'timeout' in msg: @@ -24,9 +30,13 @@ def ssh_boot_status_message(err): return str(err) -class SSHClient(): +class SSHClient: + """ + Runs commands and transfers files on a remote host over SSH. The + connection is created on first use and reused afterwards + """ - def __init__(self, ip_address, ssh_credentials): + def __init__(self, ip_address: str, ssh_credentials: Dict[str, Any]): self.ip_address = ip_address self.ssh_credentials = ssh_credentials self.ssh_client = None @@ -35,132 +45,110 @@ def __init__(self, ip_address, ssh_credentials): fpath = os.path.expanduser(self.ssh_credentials['key_filename']) self.ssh_credentials['key_filename'] = fpath if not os.path.exists(fpath): - logger.debug(f"Private key file {fpath} doesn't exist. Trying with the default key") - self.ssh_credentials['key_filename'] = os.path.expanduser('~/.ssh/id_rsa') - - def close(self): - """ - Closes the SSH client connection - """ + logger.debug( + f"Private key file {fpath} does not exist. " + "Trying with the default key" + ) + self.ssh_credentials['key_filename'] = _DEFAULT_KEY + + def close(self) -> None: + """Closes the connection, if there is one, and forgets about it""" if self.ssh_client: try: self.ssh_client.close() except Exception: + # A connection that cannot be closed is dropped anyway pass self.ssh_client = None - def create_client(self, timeout=2): - """ - Create the SSH client connection - """ - try: - self.ssh_client = paramiko.SSHClient() - self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) - - user = self.ssh_credentials.get('username') - password = self.ssh_credentials.get('password') - pkey = None - - if self.ssh_credentials.get('key_filename'): - with open(self.ssh_credentials['key_filename']) as f: - pkey = paramiko.RSAKey.from_private_key(f) + def create_client(self, timeout: int = 2) -> paramiko.SSHClient: + """Opens a new connection, replacing the current one""" + ssh_client = paramiko.SSHClient() + ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + + connect_kwargs = { + 'hostname': self.ip_address, + 'username': self.ssh_credentials.get('username'), + 'password': self.ssh_credentials.get('password'), + 'timeout': timeout, + 'banner_timeout': 200, + 'allow_agent': False, + 'look_for_keys': False, + } + key_filename = self.ssh_credentials.get('key_filename') + if key_filename: + connect_kwargs['key_filename'] = key_filename + + # Only kept once connected, so that a failed connect does not leave + # a client behind for _ensure_client to reuse + ssh_client.connect(**connect_kwargs) + self.ssh_client = ssh_client + logger.debug(f"{self.ip_address} ssh client created") + return self.ssh_client - self.ssh_client.connect( - self.ip_address, username=user, - password=password, pkey=pkey, - timeout=timeout, banner_timeout=200, - allow_agent=False, look_for_keys=False - ) + def _ensure_client(self) -> paramiko.SSHClient: + if self.ssh_client is None: + self.create_client() + return self.ssh_client - logger.debug(f"{self.ip_address} ssh client created") - except Exception as e: - raise e + @contextmanager + def _sftp(self): + ftp_client = self._ensure_client().open_sftp() + try: + yield ftp_client + finally: + ftp_client.close() - return self.ssh_client + def _exec_command(self, cmd: str, timeout: Optional[int]): + return self.ssh_client.exec_command(cmd, timeout=timeout) - def run_remote_command(self, cmd, timeout=None, run_async=False): + def run_remote_command( + self, cmd: str, timeout: Optional[int] = None, run_async: bool = False + ) -> Tuple[Optional[str], Optional[str]]: """ - Executa a command - param: timeout: execution timeout - param: run_async: do not wait for command completion + Runs a command on the remote host and returns its stdout and stderr, + or a pair of Nones when asked not to wait for it to complete """ if not self.ip_address or self.ip_address == '0.0.0.0': raise Exception('Invalid IP Address') - if self.ssh_client is None: - self.ssh_client = self.create_client() - + self._ensure_client() + # stdin is kept until this returns: closing it, which garbage + # collection does, sends EOF to the command still being read below try: - stdin, stdout, stderr = self.ssh_client.exec_command(cmd, timeout=timeout) + stdin, stdout, stderr = self._exec_command(cmd, timeout) except Exception: - # Normally this is a timeout exception - self.ssh_client = self.create_client() - stdin, stdout, stderr = self.ssh_client.exec_command(cmd, timeout=timeout) - - out = None - err = None + # The reused connection may have died since the last command + self.create_client() + stdin, stdout, stderr = self._exec_command(cmd, timeout) - if not run_async: - out = stdout.read().decode().strip() - err = stderr.read().decode().strip() - - return out, err - - def download_remote_file(self, remote_src, local_dst): - """ - Downloads a remote file to a local destination - param: local_src: local file path source - param: remote_dst: remote file path destination - """ - if self.ssh_client is None: - self.ssh_client = self.create_client() + if run_async: + return None, None + return stdout.read().decode().strip(), stderr.read().decode().strip() + def download_remote_file(self, remote_src: str, local_dst: str) -> None: + """Downloads a remote file, creating the local directory if needed""" dirname = os.path.dirname(local_dst) - if dirname and not os.path.exists(dirname): - os.makedirs(dirname) - - ftp_client = self.ssh_client.open_sftp() - ftp_client.get(remote_src, local_dst) - ftp_client.close() - - def upload_local_file(self, local_src, remote_dst): - """ - Upload a local file to a rempote destination - param: local_src: local file path source - param: remote_dst: remote file path destination - """ - if self.ssh_client is None: - self.ssh_client = self.create_client() - - ftp_client = self.ssh_client.open_sftp() - ftp_client.put(local_src, remote_dst) - ftp_client.close() - - def upload_multiple_local_files(self, file_list): - """ - upload multiple files with the same sftp connection - param: file_list: list of tuples [(local_src, remote_dst),] - """ - if self.ssh_client is None: - self.ssh_client = self.create_client() - - ftp_client = self.ssh_client.open_sftp() - for local_src, remote_dst in file_list: + if dirname: + os.makedirs(dirname, exist_ok=True) + with self._sftp() as ftp_client: + ftp_client.get(remote_src, local_dst) + + def upload_local_file(self, local_src: str, remote_dst: str) -> None: + """Uploads a local file to a remote destination""" + with self._sftp() as ftp_client: ftp_client.put(local_src, remote_dst) - ftp_client.close() - - def upload_data_to_file(self, data, remote_dst): - """ - upload data to a remote file - param: data: string data - param: remote_dst: remote file path destination - """ - if self.ssh_client is None: - self.ssh_client = self.create_client() - - ftp_client = self.ssh_client.open_sftp() - - with ftp_client.open(remote_dst, 'w') as f: - f.write(data) - ftp_client.close() + def upload_multiple_local_files( + self, file_list: List[Tuple[str, str]] + ) -> None: + """Uploads several local files reusing a single SFTP connection""" + with self._sftp() as ftp_client: + for local_src, remote_dst in file_list: + ftp_client.put(local_src, remote_dst) + + def upload_data_to_file(self, data: str, remote_dst: str) -> None: + """Writes data into a remote file""" + with self._sftp() as ftp_client: + with ftp_client.open(remote_dst, 'w') as remote_file: + remote_file.write(data) diff --git a/lithops/utils.py b/lithops/utils.py index f7ac3d4bf..167e61a27 100644 --- a/lithops/utils.py +++ b/lithops/utils.py @@ -35,6 +35,8 @@ from enum import Enum from contextlib import closing +from typing import List + from lithops import constants from lithops.version import __version__ @@ -42,12 +44,46 @@ logger = logging.getLogger(__name__) +class ShutdownSafeStreamHandler(logging.StreamHandler): + """StreamHandler that does not traceback when the stream is already closed.""" + + def emit(self, record): + try: + stream = self.stream + if stream is None or getattr(stream, 'closed', False): + return + msg = self.format(record) + stream.write(msg + self.terminator) + self.flush() + except RecursionError: + # handleError() logs, so it would recurse again. Same as logging + raise + except (ValueError, OSError): + # The stream was closed between the check above and the write + return + except Exception: + self.handleError(record) + + def uuid_str(): return str(uuid.uuid4()) +def _as_future_list(fs): + """Wrap a single future; leave list / FuturesList unchanged.""" + return fs if isinstance(fs, list) else [fs] + + +def _future_id(fut): + return (fut.executor_id, fut.job_id, fut.call_id) + + def create_executor_id(lenght=6): - """ Creates an executor ID. """ + """ + Creates the ID of a new executor. Executors of the same session share the + session ID and are told apart by a counter, both kept in the environment + so that they survive across processes + """ if '__LITHOPS_SESSION_ID' in os.environ: session_id = os.environ['__LITHOPS_SESSION_ID'] else: @@ -60,14 +96,53 @@ def create_executor_id(lenght=6): exec_num = 0 os.environ['__LITHOPS_TOTAL_EXECUTORS'] = str(exec_num) - return '{}-{}'.format(session_id, exec_num) + return f'{session_id}-{exec_num}' + + +# Carries the monitoring queues of an executor down to the workers, so that an +# executor created inside one of them can extend the chain +MONITORING_QUEUES_ENV = '__LITHOPS_MONITORING_QUEUES' + + +def monitoring_queue_name(executor_id: str) -> str: + """Returns the name of the queue an executor is monitored through""" + return f'lithops-{executor_id}' + + +def monitoring_queues(executor_id: str) -> List[str]: + """ + Returns every queue a call status of this executor has to be published to: + the queue of each executor up the chain, ending with this one. + + The chain travels in the environment instead of being read back out of the + executor id, because the id does not say how deep it is: a worker adds the + job and the call to the session id while a remote invoker adds only the + job, so the same number of tokens can stand for different chains + """ + parent_queues = [] + raw_queues = os.environ.get(MONITORING_QUEUES_ENV) + if raw_queues: + try: + parent_queues = list(json.loads(raw_queues)) + except ValueError: + logger.warning( + f'Ignoring a malformed {MONITORING_QUEUES_ENV}: {raw_queues}' + ) + + queue = monitoring_queue_name(executor_id) + if queue in parent_queues: + # The remote invoker builds the payload of the job it spawns from + # inside a worker that already exported this very chain, so without + # this every remotely invoked task would report twice to the client + return parent_queues + return parent_queues + [queue] def get_executor_id(): - """ retrieves the current executor ID. """ + """Returns the ID of the last executor created in this session""" session_id = os.environ['__LITHOPS_SESSION_ID'] exec_num = os.environ['__LITHOPS_TOTAL_EXECUTORS'] - return '{}-{}'.format(session_id, exec_num) + return f'{session_id}-{exec_num}' def iterchunks(lst, n): @@ -77,8 +152,9 @@ def iterchunks(lst, n): def agg_data(data_strs): - """Auxiliary function that aggregates data of a job to a single - byte string. + """ + Concatenates the data of every call of a job into a single byte string, + and returns it along with the byte range that each call occupies """ ranges = [] pos = 0 @@ -90,7 +166,7 @@ def agg_data(data_strs): def create_futures_list(futures, executor): - """creates a new FuturesList an initiates its attrs""" + """Creates a new FuturesList bound to the executor that produced it""" fl = FuturesList(futures) fl.config = executor.config fl.executor = executor @@ -99,13 +175,27 @@ def create_futures_list(futures, executor): class FuturesList(list): + """ + List of futures that can be mapped over again, so that jobs can be + chained. Chaining replaces the contents with the futures of the new job, + while alt_list keeps every future of the chain for wait() and get_result() + """ + + # Defaults for lists that were not built by create_futures_list, and for + # the ones rehydrated by __reduce__, which drops the executor + executor = None + config = None def _create_executor(self): if not self.executor: from lithops import FunctionExecutor self.executor = FunctionExecutor(config=self.config) + def _all_futures(self): + return self.alt_list if hasattr(self, 'alt_list') else self + def _extend_futures(self, fs): + # Only the last job of the chain produces the output of the chain for fut in self: fut._produce_output = False if not hasattr(self, 'alt_list'): @@ -127,67 +217,85 @@ def map_reduce(self, map_function, reduce_function, sync=False, **kwargs): self._create_executor() if sync: self.executor.wait(self) - fs = self.executor.map_reduce(map_function, self, reduce_function, **kwargs) + fs = self.executor.map_reduce( + map_function, self, reduce_function, **kwargs + ) self._extend_futures(fs) return self def wait(self, **kwargs): self._create_executor() - fs_tt = self.alt_list if hasattr(self, 'alt_list') else self - return self.executor.wait(fs_tt, **kwargs) + return self.executor.wait(self._all_futures(), **kwargs) def get_result(self, **kwargs): self._create_executor() - fs_tt = self.alt_list if hasattr(self, 'alt_list') else self - return self.executor.get_result(fs_tt, **kwargs) + return self.executor.get_result(self._all_futures(), **kwargs) def __reduce__(self): + # The executor is not picklable, and a rehydrated list creates its own self.executor = None return super().__reduce__() -def get_default_backend(mode): - """ Return lithops execution backend """ +_MODE_TO_DEFAULT_BACKEND = { + constants.LOCALHOST: constants.LOCALHOST, + constants.SERVERLESS: constants.SERVERLESS_BACKEND_DEFAULT, + constants.STANDALONE: constants.STANDALONE_BACKEND_DEFAULT, +} - if mode == constants.LOCALHOST: - return constants.LOCALHOST - elif mode == constants.SERVERLESS: - return constants.SERVERLESS_BACKEND_DEFAULT - elif mode == constants.STANDALONE: - return constants.STANDALONE_BACKEND_DEFAULT - elif mode: - raise Exception("Unknown exeution mode: {}".format(mode)) + +def get_default_backend(mode): + """Returns the compute backend an execution mode defaults to""" + if mode in _MODE_TO_DEFAULT_BACKEND: + return _MODE_TO_DEFAULT_BACKEND[mode] + if mode: + raise Exception(f"Unknown execution mode: {mode}") def get_mode(backend): - """ Return lithops execution mode """ - + """Returns the execution mode a compute backend belongs to""" if backend is None: return constants.MODE_DEFAULT if backend == constants.LOCALHOST: return constants.LOCALHOST - elif backend in constants.SERVERLESS_BACKENDS: + if backend in constants.SERVERLESS_BACKENDS: return constants.SERVERLESS - elif backend in constants.STANDALONE_BACKENDS: + if backend in constants.STANDALONE_BACKENDS: return constants.STANDALONE - elif backend: - raise Exception("Unknown compute backend: {}".format(backend)) + if backend: + raise Exception(f"Unknown compute backend: {backend}") + + +def log_prefix(executor_id, job_id=None, call_id=None) -> str: + """Identity prefix used in Lithops log messages""" + parts = [f'ExecutorID {executor_id}'] + if job_id is not None: + parts.append(f'JobID {job_id}') + if call_id is not None: + parts.append(f'CallID {call_id}') + return ' | '.join(parts) def setup_lithops_logger(log_level=constants.LOGGER_LEVEL, log_format=constants.LOGGER_FORMAT, stream=None, filename=None): - """Setup logging for lithops.""" + """ + Configures the lithops logger. A log level of None, or 'none', leaves the + logging of the process untouched + """ if log_level is None or str(log_level).lower() == 'none': return if stream is None: stream = constants.LOGGER_STREAM + # Both handlers are always declared, so the unused FileHandler is pointed + # at os.devnull rather than at a file nobody asked for + log_to_file = filename is not None if filename is None: filename = os.devnull - if type(log_level) is str: + if isinstance(log_level, str): log_level = logging.getLevelName(log_level.upper()) config_dict = { @@ -202,7 +310,7 @@ def setup_lithops_logger(log_level=constants.LOGGER_LEVEL, 'console_handler': { 'level': log_level, 'formatter': 'standard', - 'class': 'logging.StreamHandler', + 'class': 'lithops.utils.ShutdownSafeStreamHandler', 'stream': stream }, 'file_handler': { @@ -222,116 +330,151 @@ def setup_lithops_logger(log_level=constants.LOGGER_LEVEL, } } - if filename is not os.devnull: + if log_to_file: config_dict['loggers']['lithops']['handlers'] = ['file_handler'] logging.config.dictConfig(config_dict) -def create_handler_zip(dst_zip_location, entry_point_files, entry_point_name=None): - """Create the zip package that is uploaded as a function""" +_SKIP_HANDLER_ZIP_DIRS = frozenset({'__pycache__', '.pytest_cache'}) + + +def _skip_in_handler_zip(path: str, dst_zip_location: str) -> bool: + # The zip is often written inside the package directory that is being + # zipped, so it must not add itself, nor any other package left there + return os.path.abspath(path) == dst_zip_location or path.endswith('.zip') + + +def _add_folder_to_handler_zip( + zip_file: zipfile.ZipFile, + full_dir_path: str, + dst_zip_location: str, + sub_dir: str = '' +) -> None: + """Adds a directory tree to the zip, under the lithops/ prefix""" + for name in os.listdir(full_dir_path): + full_path = os.path.join(full_dir_path, name) + if os.path.isdir(full_path): + if name not in _SKIP_HANDLER_ZIP_DIRS: + _add_folder_to_handler_zip( + zip_file, + full_path, + dst_zip_location, + os.path.join(sub_dir, name), + ) + elif os.path.isfile(full_path): + if not _skip_in_handler_zip(full_path, dst_zip_location): + zip_file.write( + full_path, os.path.join('lithops', sub_dir, name) + ) + - logger.debug("Creating function handler zip in {}".format(dst_zip_location)) +def create_handler_zip( + dst_zip_location, entry_point_files, entry_point_name=None +): + """ + Creates the zip package that is uploaded as a function: the entry points + at its root, and the whole lithops package under lithops/ + """ + dst_zip_location = os.path.abspath(dst_zip_location) + logger.debug(f"Creating function handler zip in {dst_zip_location}") - def add_folder_to_zip(zip_file, full_dir_path, sub_dir=''): - for file in os.listdir(full_dir_path): - full_path = os.path.join(full_dir_path, file) - if os.path.isfile(full_path): - zip_file.write(full_path, os.path.join('lithops', sub_dir, file)) - elif os.path.isdir(full_path) and '__pycache__' not in full_path: - add_folder_to_zip(zip_file, full_path, os.path.join(sub_dir, file)) + if not isinstance(entry_point_files, list): + entry_point_files = [entry_point_files] + created = False try: - ep_files = entry_point_files if isinstance(entry_point_files, list) else [entry_point_files] - with zipfile.ZipFile(dst_zip_location, 'w', zipfile.ZIP_DEFLATED) as lithops_zip: - module_location = os.path.dirname(os.path.abspath(lithops.__file__)) - for ep_file in ep_files: + with zipfile.ZipFile( + dst_zip_location, 'w', zipfile.ZIP_DEFLATED + ) as lithops_zip: + module_location = os.path.dirname( + os.path.abspath(lithops.__file__) + ) + for ep_file in entry_point_files: ep_name = entry_point_name or os.path.basename(ep_file) lithops_zip.write(ep_file, ep_name) - add_folder_to_zip(lithops_zip, module_location) - + _add_folder_to_handler_zip( + lithops_zip, module_location, dst_zip_location + ) + created = True + zip_size = os.path.getsize(dst_zip_location) + logger.debug( + f'Function handler zip created - Size: {sizeof_fmt(zip_size)}' + ) except Exception as e: - raise Exception(f'Unable to create the {dst_zip_location} package: {e}') + raise Exception( + f'Unable to create the {dst_zip_location} package: {e}' + ) from e + finally: + # A half written zip would be uploaded and fail at invocation time + if not created and os.path.exists(dst_zip_location): + os.remove(dst_zip_location) -def verify_runtime_name(runtime_name): - """Check if the runtime name has a correct formating""" +def verify_runtime_name(runtime_name: str) -> None: + """Asserts that the runtime name can be used as a container image name""" assert re.match("^[A-Za-z0-9_/.:-]*$", runtime_name), \ f'Runtime name "{runtime_name}" not valid' def timeout_handler(error_msg, signum, frame): + """Signal handler that turns an alarm into a TimeoutError""" raise TimeoutError(error_msg) -def version_str(version_info): - """Format the python version information""" - return "{}.{}".format(version_info[0], version_info[1]) +def version_str(version_info) -> str: + """Formats a sys.version_info tuple as major.minor""" + return f"{version_info[0]}.{version_info[1]}" -def is_unix_system(): - """Check if the current OS is UNIX""" - curret_system = platform.system() - return curret_system != 'Windows' +def is_unix_system() -> bool: + """Checks if the current OS is UNIX""" + return platform.system() != 'Windows' -def is_linux_system(): - """Check if the current OS is LINUX""" - curret_system = platform.system().lower() - if curret_system == "linux": - return True - else: - return False +def is_linux_system() -> bool: + """Checks if the current OS is LINUX""" + return platform.system().lower() == "linux" -def is_lithops_worker(): - """ - Checks if the current execution is within a lithops worker - """ - if 'LITHOPS_WORKER' in os.environ: - return True - return False +def is_lithops_worker() -> bool: + """Checks if the current execution is within a lithops worker""" + return 'LITHOPS_WORKER' in os.environ -def is_object_processing_function(map_function): +def is_object_processing_function(map_function) -> bool: """ Checks if a function contains the obj parameter, which means - the user wants to activate the data processing logic. + the user wants to activate the data processing logic """ func_sig = inspect.signature(map_function) - return {'obj'} & set(func_sig.parameters) + return 'obj' in func_sig.parameters -def is_notebook(): +def is_notebook() -> bool: + """Checks if the current execution is within a Jupyter notebook""" try: - shell = get_ipython().__class__.__name__ - if shell == 'ZMQInteractiveShell': - return True # Jupyter notebook or qtconsole - elif shell == 'TerminalInteractiveShell': - return False # Terminal running IPython - else: - return False # Other type (?) + return get_ipython().__class__.__name__ == 'ZMQInteractiveShell' except NameError: - return False # Probably standard Python interpreter + return False def convert_bools_to_string(extra_env): - """ - Converts all booleans of a dictionary to a string - """ - for key in extra_env: - if type(extra_env[key]) is bool: - extra_env[key] = str(extra_env[key]) + """Converts every boolean value of a dictionary to a string, in place""" + for key, value in extra_env.items(): + if isinstance(value, bool): + extra_env[key] = str(value) return extra_env -def sizeof_fmt(num, suffix='B'): +def sizeof_fmt(num, suffix='B') -> str: + """Formats a number of bytes with a binary unit prefix""" for unit in ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi']: if abs(num) < 1024.0: - return "%3.1f%s%s" % (num, unit, suffix) + return f'{num:3.1f}{unit}{suffix}' num /= 1024.0 - return "%.1f%s%s" % (num, 'Yi', suffix) + return f'{num:.1f}Yi{suffix}' def sdb_to_dict(item): @@ -364,7 +507,8 @@ def b64str_to_bytes(str_data): return byte_data -def get_docker_path(): +def get_docker_path() -> str: + """Returns the path of the docker command, or of podman as a fallback""" docker_path = shutil.which('docker') podman_path = shutil.which('podman') if not docker_path and not podman_path: @@ -373,112 +517,134 @@ def get_docker_path(): return docker_path or podman_path +def _get_required_param(backend_config, backend: str, param: str): + if param not in backend_config: + raise Exception( + f'You must provide "{param}" param in config ' + f'under "{backend}" section' + ) + return backend_config[param] + + def get_default_container_name(backend, backend_config, runtime_name): """ - Generates the default runtime image name - Used in serverless/kubernetes-based backends + Generates the default runtime image name, qualified with the registry the + backend is configured to use. Used in serverless and kubernetes backends """ python_version = CURRENT_PY_VERSION.replace('.', '') img = f'{runtime_name}-v{python_version}:{__version__}' docker_server = backend_config['docker_server'] + # Every registry qualifies the image with a different set of params, so + # they are recognised by their well known hostnames if 'docker.io' in docker_server: # Docker hub container registry - try: - docker_user = backend_config['docker_user'] - except Exception: - raise Exception('You must provide "docker_user" param ' - f'in config under "{backend}" section') + docker_user = _get_required_param( + backend_config, backend, 'docker_user' + ) return f'docker.io/{docker_user}/{img}' elif 'icr.io' in docker_server: # IBM container registry - try: - docker_namespace = backend_config['docker_namespace'] - except Exception: - raise Exception('You must provide "docker_namespace" param' - f'in config under "{backend}" section') + docker_namespace = _get_required_param( + backend_config, backend, 'docker_namespace' + ) return f'{docker_server}/{docker_namespace}/{img}' elif 'pkg.dev' in docker_server: # Google Artifact Registry (Docker) - try: - region = backend_config['region'] - project_name = backend_config['project_name'] - repository = backend_config.get('artifact_registry_repository', 'lithops') - except Exception: - raise Exception('You must provide "region" and "project_name" params' - 'in config under "gcp" section') + if 'region' not in backend_config or 'project_name' not in backend_config: + raise Exception( + 'You must provide "region" and "project_name" params in ' + 'config under "gcp" section' + ) + region = backend_config['region'] + project_name = backend_config['project_name'] + repository = backend_config.get( + 'artifact_registry_repository', 'lithops' + ) return f'{region}-docker.pkg.dev/{project_name}/{repository}/{img}' else: return f'{docker_server}/{img}' +def _get_docker_desktop_username() -> str: + """Reads the registry user out of the Docker Desktop credential helper""" + cmd = ( + "docker-credential-desktop list | jq -r 'to_entries[].key' | while " + "read; do docker-credential-desktop get <<<$REPLY; break; done" + ) + try: + credentials = sp.check_output( + cmd, shell=True, encoding='UTF-8', stderr=sp.STDOUT + ) + return json.loads(credentials)['Username'] + except Exception: + raise Exception('Unable to get the Docker registry user') + + def get_docker_username(): - user = None + """Returns the user that docker/podman is logged in to the registry as""" docker_path = get_docker_path() - - docker_user_info = sp.check_output( + docker_info = sp.check_output( f"{docker_path} info", shell=True, encoding='UTF-8', stderr=sp.STDOUT ) - for line in docker_user_info.splitlines(): - if 'Username' in line: - _, useranme = line.strip().split(':') - user = useranme.strip() - if user is None: - try: - cmd = ("docker-credential-desktop list | jq -r 'to_entries[].key' | while " - "read; do docker-credential-desktop get <<<$REPLY; break; done") - docker_user_info = sp.check_output(cmd, shell=True, encoding='UTF-8', stderr=sp.STDOUT) - docker_data = json.loads(docker_user_info) - user = docker_data['Username'] - except Exception: - raise Exception('Unable to get the Docker registry user') + user = None + for line in docker_info.splitlines(): + if 'Username' in line: + _, username = line.strip().split(':') + user = username.strip() - return user + return user if user is not None else _get_docker_desktop_username() -def find_free_port(): +def find_free_port() -> int: + """Returns a port that is free at this instant""" with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s: - s.bind(('', 0)) s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + s.bind(('', 0)) return s.getsockname()[1] +# Schemes that name a Lithops storage backend by its protocol +_URL_SCHEME_ALIASES = {'cos': 'ibm_cos', 's3': 'aws_s3'} + + def split_object_url(obj_url): + """ + Splits a data URL into its storage backend, bucket, prefix and object + name. A URL ending in '/' names a folder, so it has no object name + """ if '://' in obj_url: sb, path = obj_url.split('://') else: sb = None path = obj_url - sb = 'ibm_cos' if sb == 'cos' else sb - sb = 'aws_s3' if sb == 's3' else sb + sb = _URL_SCHEME_ALIASES.get(sb, sb) bucket, full_key = path.split('/', 1) if '/' in path else (path, '') if full_key.endswith('/'): - prefix = ''.join(full_key.rsplit('/', 1)) - obj_name = '' - elif full_key: - prefix, obj_name = full_key.rsplit('/', 1) if '/' in full_key else ('', full_key) + prefix, obj_name = full_key[:-1], '' + elif '/' in full_key: + prefix, obj_name = full_key.rsplit('/', 1) else: - prefix = '' - obj_name = '' + prefix, obj_name = '', full_key return sb, bucket, prefix, obj_name def split_path(path): - - if (path.startswith("/")): + """Splits a storage path into its bucket and its key""" + if path.startswith("/"): path = path[1:] ind = path.find("/") - if (ind > 0): + if ind > 0: bucket_name = path[:ind] key = path[ind + 1:] else: @@ -489,106 +655,108 @@ def split_path(path): def format_data(iterdata, extra_args): """ - Converts iteradata to a list with extra_args + Converts iterdata to a list, appending extra_args to every element. The + element decides how: tuples are concatenated, dicts are merged """ - # Format iterdata in a proper way - if type(iterdata) in [range, set]: - data = list(iterdata) - elif type(iterdata) is not list and type(iterdata) is not FuturesList: - data = [iterdata] - else: - data = iterdata - - if extra_args: - new_iterdata = [] - for data_i in data: - - if type(data_i) is tuple: - # multiple args - if type(extra_args) is not tuple: - raise Exception('extra_args must contain args in a tuple') - new_iterdata.append(data_i + extra_args) - - elif type(data_i) is dict: - # kwargs - if type(extra_args) is not dict: - raise Exception('extra_args must contain kwargs in a dictionary') - data_i.update(extra_args) - new_iterdata.append(data_i) - else: - new_iterdata.append((data_i, *extra_args)) - data = new_iterdata + data = _as_iterdata_list(iterdata) + if not extra_args: + return data + + new_iterdata = [] + for data_i in data: + if type(data_i) is tuple: + if type(extra_args) is not tuple: + raise Exception( + 'extra_args must contain args in a tuple' + ) + new_iterdata.append(data_i + extra_args) + elif type(data_i) is dict: + if type(extra_args) is not dict: + raise Exception( + 'extra_args must contain kwargs in a dictionary' + ) + data_i.update(extra_args) + new_iterdata.append(data_i) + else: + new_iterdata.append((data_i, *extra_args)) + return new_iterdata - return data +def _as_iterdata_list(iterdata): + if isinstance(iterdata, (range, set)): + return list(iterdata) + if isinstance(iterdata, list): + return iterdata + return [iterdata] -def verify_args(func, iterdata, extra_args): - if isinstance(iterdata, FuturesList): - # this is required for function chaining - return [{'future': f} for f in iterdata] +# Params that Lithops injects at invocation time, so the user is not expected +# to provide them in the iterdata +_INJECTED_ARGS = frozenset({'ibm_cos', 'storage', 'id', 'rabbitmq'}) - data = format_data(iterdata, extra_args) - # Verify parameters - non_verify_args = ['ibm_cos', 'storage', 'id', 'rabbitmq'] +def _user_signature(func) -> inspect.Signature: + """Signature of a map function, without the params Lithops injects""" func_sig = inspect.signature(func) - - new_parameters = [ + user_parameters = [ param for name, param in func_sig.parameters.items() - if name not in non_verify_args + if name not in _INJECTED_ARGS ] - new_func_sig = func_sig.replace(parameters=new_parameters) + return func_sig.replace(parameters=user_parameters) + + +def verify_args(func, iterdata, extra_args): + """ + Binds every element of the iterdata to the params of the map function, + returning one kwargs dict per call + """ + if isinstance(iterdata, FuturesList): + # A chained job receives the future of the previous one, which is only + # bound to a param once the previous job finishes + return [{'future': f} for f in iterdata] + + data = format_data(iterdata, extra_args) + func_sig = _user_signature(func) - # Detect presence of **kwargs (with any name) + # A wrapper, such as a decorator, hides the params of the function behind + # **kwargs, so the names of a dict element cannot be checked against them has_var_keyword = any( p.kind == inspect.Parameter.VAR_KEYWORD - for p in new_func_sig.parameters.values() + for p in func_sig.parameters.values() ) new_data = [] for elem in data: if isinstance(elem, dict): - # If the function accepts **kwargs (any name), we cannot reliably - # enforce exact param name matching here, and we *want* to allow - # passing through arbitrary dicts (e.g., original function args) - # even when a decorator wrapper has **kwargs, etc. - if has_var_keyword: - new_data.append(elem) - elif set(expected_keys := list(new_func_sig.parameters)) <= set(elem): - # No **kwargs: enforce that the dict contains at least all - # required user parameters (excluding reserved ones). + if has_var_keyword or set(func_sig.parameters) <= set(elem): new_data.append(elem) else: raise ValueError( - "Check the args names in the data. You provided these args: ", - f"{list(elem)}, and the args must be: {expected_keys}", + "Check the args names in the data. You provided these " + f"args: {list(elem)}, and the args must be: " + f"{list(func_sig.parameters)}" ) elif isinstance(elem, tuple): - new_elem = dict(new_func_sig.bind(*elem).arguments) - new_data.append(new_elem) + new_data.append(dict(func_sig.bind(*elem).arguments)) else: - # single value (list, string, integer, dict, etc) - new_elem = dict(new_func_sig.bind(elem).arguments) - new_data.append(new_elem) + # A single value of any other type binds to the first param + new_data.append(dict(func_sig.bind(elem).arguments)) return new_data class WrappedStreamingBody: """ - Wrap boto3's StreamingBody object to provide enough Python fileobj functionality. + Wrap boto3's StreamingBody object to provide enough + Python fileobj functionality. from https://gist.github.com/debedb/2e5cbeb54e43f031eaf0 """ def __init__(self, sb, size): - # The StreamingBody we're wrapping self.sb = sb - # Initial position self.pos = 0 - # Size of the object self.size = size def tell(self): @@ -622,7 +790,6 @@ def seek(self, offset, whence=0): retval = self.size else: retval = offset - # print("In seek(%s, %s): %s, size is %s" % (offset, whence, retval, self.size)) self.pos = retval return retval @@ -637,49 +804,32 @@ def __next__(self): return self.read(64 * 1024) def __getattr__(self, attr): - if attr == 'tell': - return self.tell - elif attr == 'seek': - return self.seek - elif attr == 'read': - return self.read - elif attr == 'readline': - return self.readline - elif attr == '__str__': - return self.__str__ - elif attr == '__iter__': - return self.__iter__ - elif attr == '__next__': - return self.__next__ - else: - return getattr(self.sb, attr) + # Only reached for the attributes this wrapper does not define, so + # everything else of the fileobj protocol falls through to boto3 + return getattr(self.sb, attr) class WrappedStreamingBodyPartition(WrappedStreamingBody): """ - Wrap boto3's StreamingBody object to provide line integrity of the partitions - based on the newline character. + Wrap boto3's StreamingBody object to provide line + integrity of the partitions based on the newline + character. """ def __init__(self, sb, size, byterange, newline='\n'): super().__init__(sb, size) - # Range of the chunk self.range = byterange - # New line character self.newline_char = newline.encode() - # The first chunk does not contain plusbyte + # Every chunk but the first one reads one byte early, so that read() + # can tell whether the previous chunk ended in the middle of a row self._plusbytes = 0 if not self.range or self.range[0] == 0 else 1 - # To store the first byte of this chunk, which actually is the last byte of previous chunk self._first_byte = None - # Flag that indicates the end of the file self._eof = False - # special logic the first time the stream is read self._first_read = True def read(self, n=None): if self._eof: return b'' - # Data always contain one byte from the previous chunk, - # so l'ets check if it is a \n or not + if not self._first_byte and self._plusbytes == 1: self._first_byte = self.sb.read(self._plusbytes) @@ -690,13 +840,13 @@ def read(self, n=None): if self._first_read and self._first_byte and \ self._first_byte != self.newline_char: + # The previous chunk did not end in a newline, so the first row of + # this one is a cut row that the previous chunk already returned logger.debug('Discarding first partial row') - # Previous byte is not self.newline_char - # This means that we have to discard first row because it is cut first_row_start_pos = retval.find(self.newline_char) + 1 self._first_read = False - # Find end of the line in threshold + # The last row of a chunk is completed past its own end if self.pos >= self.size: current_end_pos = last_row_end_pos - (self.pos - self.size) last_byte_pos = retval[current_end_pos - 1:].find(self.newline_char) @@ -752,24 +902,31 @@ def docker_login(docker_user, docker_password, docker_server): def run_command(cmd, return_result=False, input=None): - kwargs = {} - - if logger.getEffectiveLevel() != logging.DEBUG: - kwargs['stderr'] = sp.DEVNULL + """ + Runs a shell command, silencing its output unless lithops is in debug + mode. Returns its stdout if asked to, otherwise nothing + """ + quiet = logger.getEffectiveLevel() != logging.DEBUG + kwargs = {'stderr': sp.DEVNULL} if quiet else {} if input: - return sp.check_output(cmd.split(), input=bytes(input, 'utf-8'), **kwargs) + return sp.check_output( + cmd.split(), + input=bytes(input, 'utf-8'), + **kwargs, + ) if return_result: result = sp.check_output(cmd.split(), encoding='UTF-8', **kwargs) return result.strip().replace('"', '') - else: - if logger.getEffectiveLevel() != logging.DEBUG: - kwargs['stdout'] = sp.DEVNULL - sp.check_call(cmd.split(), **kwargs) + + if quiet: + kwargs['stdout'] = sp.DEVNULL + sp.check_call(cmd.split(), **kwargs) -def is_podman(docker_path): +def is_podman(docker_path) -> bool: + """Checks whether the docker command is actually podman""" try: cmd = f'{docker_path} info | grep podman' sp.check_output(cmd, shell=True, stderr=sp.STDOUT) @@ -784,6 +941,11 @@ class BackendType(Enum): class CountDownLatch: + """ + Barrier that blocks the waiters until it has been unlocked as many times + as the count it was created with + """ + def __init__(self, count): self.count = count self.event = threading.Event() diff --git a/lithops/wait.py b/lithops/wait.py index 22f12b2e8..c58f31949 100644 --- a/lithops/wait.py +++ b/lithops/wait.py @@ -24,8 +24,16 @@ from itertools import chain from typing import Optional, List, Union, Tuple, Any -from lithops.utils import is_unix_system, timeout_handler, \ - is_notebook, is_lithops_worker, FuturesList +from lithops.utils import ( + is_unix_system, + timeout_handler, + is_notebook, + is_lithops_worker, + FuturesList, + _as_future_list, + _future_id, + log_prefix, +) from lithops.storage import InternalStorage from lithops.future import ResponseFuture from lithops.monitor import JobMonitor @@ -41,35 +49,165 @@ logger = logging.getLogger(__name__) -def wait(fs: Union[ResponseFuture, FuturesList, List[ResponseFuture]], - internal_storage: Optional[InternalStorage] = None, - job_monitor: Optional[JobMonitor] = None, - throw_except: Optional[bool] = True, - return_when: Optional[Any] = ALL_COMPLETED, - download_results: Optional[bool] = False, - timeout: Optional[int] = None, - threadpool_size: Optional[int] = THREADPOOL_SIZE, - wait_dur_sec: Optional[int] = None, - show_progressbar: Optional[bool] = True, - futures_from_executor_wait: Optional[bool] = False) -> Tuple[FuturesList, FuturesList]: +def _future_is_complete(fut, download_results): + return fut.done if download_results else (fut.success or fut.done) + + +def _partition_futures(fs, download_results): + """Split futures into (done, not_done).""" + done, not_done = [], [] + for fut in fs: + if _future_is_complete(fut, download_results): + done.append(fut) + else: + not_done.append(fut) + return done, not_done + + +def _poll_sleep_sec(job_monitor, wait_dur_sec): + """ + Returns the interval between two polls. Only a remote storage backend is + worth throttling, every other source of statuses is local and cheap + """ + remote_storage = ( + job_monitor.type == 'storage' + and job_monitor.storage_backend != 'localhost' + ) + if remote_storage: + return wait_dur_sec or WAIT_DUR_SEC + return 0.1 + + +def _log_wait_start(prefix: str, return_when: Any, pending: int) -> None: + """ + Logs how many function activations the wait is going to block on """ - Wait for the Future instances (possibly created by different Executor instances) - given by fs to complete. Returns a named 2-tuple of sets. The first set, named done, - contains the futures that completed (finished or cancelled futures) before the wait - completed. The second set, named not_done, contains the futures that did not complete - (pending or running futures). timeout can be used to control the maximum number of - seconds to wait before returning. + if return_when == ALL_COMPLETED: + target = '' + elif return_when == ANY_COMPLETED: + target = 'any of ' + else: + target = f'{return_when}% of ' + + logger.info( + f'{prefix} - Waiting for {target}{pending} ' + 'function activations to complete' + ) + + +def _set_wait_alarm(timeout: int) -> None: + """ + Arms a SIGALRM that aborts the wait once the timeout is exceeded + """ + logger.debug(f'Setting waiting timeout to {timeout} seconds') + error_msg = ( + f'Timeout of {timeout} seconds exceeded waiting for ' + 'function activations to finish' + ) + signal.signal(signal.SIGALRM, partial(timeout_handler, error_msg)) + signal.alarm(timeout) + + +def _create_progressbar(total: int, initial: int): + """ + Builds the bar that tracks the wait, or nothing where a bar would only + get in the way: inside a worker, or when debug logs are already printed + """ + if is_lithops_worker() or logger.getEffectiveLevel() == logging.DEBUG: + return None + + from tqdm.auto import tqdm + if not is_notebook(): + print() + pbar = tqdm( + bar_format=' {l_bar}{bar}| {n_fmt}/{total_fmt} ', + total=total, + disable=None, + ) + pbar.update(min(initial, total)) + return pbar + + +def _start_job_monitors(executors_data) -> List[JobMonitor]: + """ + Starts one monitor per executor the futures belong to + """ + monitors = [] + for executor_data in executors_data: + monitor = JobMonitor( + executor_id=executor_data.executor_id, + internal_storage=executor_data.internal_storage, + ) + monitor.start(fs=executor_data.futures) + monitors.append(monitor) + return monitors + + +def _poll_until_done( + fs, + executors_data, + job_monitor, + return_when, + download_results, + sleep_sec, + poll_kwargs, +): + """ + Polls every executor until return_when% of the futures are done. A round + that fetched something is followed immediately by another one, as more + statuses are likely to be waiting already + """ + while not _check_done(fs, return_when, download_results): + # The monitor is a daemon thread that exits on its own once every + # future it knows about is done, so it may need waking up for the + # futures that showed up afterwards + if not job_monitor.is_alive(): + job_monitor.start(fs=fs) + + new_data = False + for executor_data in executors_data: + if _get_executor_data(fs, executor_data, **poll_kwargs): + new_data = True + + time.sleep(0 if new_data else sleep_sec) + + +def wait( + fs: Union[ResponseFuture, FuturesList, List[ResponseFuture]], + internal_storage: Optional[InternalStorage] = None, + job_monitor: Optional[JobMonitor] = None, + throw_except: Optional[bool] = True, + return_when: Optional[Any] = ALL_COMPLETED, + download_results: Optional[bool] = False, + timeout: Optional[int] = None, + threadpool_size: Optional[int] = THREADPOOL_SIZE, + wait_dur_sec: Optional[int] = None, + show_progressbar: Optional[bool] = True, + futures_from_executor_wait: Optional[bool] = False, +) -> Tuple[FuturesList, FuturesList]: + """ + Wait for the Future instances (possibly created by different + Executor instances) given by fs to complete. Returns a 2-tuple. + The first item, done, contains the futures that completed + before the wait completed. The second item, not_done, contains + the futures that did not complete. timeout can be used to + control the maximum number of seconds to wait before returning. :param fs: Futures list. Default None :param internal_storage: InternalStorage instance. Default None. :param job_monitor: JobMonitor instance. Default None. - :param throw_except: Re-raise exception if call raised. Default True. + :param throw_except: Re-raise exception if call raised. + Default True. :param return_when: Percentage of done futures - :param download_results: Download results. Default false (Only get statuses) + :param download_results: Download results. Default false + (Only get statuses) :param timeout: Timeout of waiting for results. :param threadpool_size: Number of threads to use. Default 64 - :param wait_dur_sec: Time interval between each check. Default 1 second + :param wait_dur_sec: Time interval between each check. + Default 1 second :param show_progressbar: whether or not to show the progress bar. + :param futures_from_executor_wait: Measure progress against the + futures that are still pending instead of against all of them. :return: `(fs_done, fs_notdone)` where `fs_done` is a list of futures that have completed @@ -77,96 +215,73 @@ def wait(fs: Union[ResponseFuture, FuturesList, List[ResponseFuture]], :rtype: 2-tuple of list """ if not fs: - return + return [], [] - if type(fs) is not list and type(fs) is not FuturesList: - fs = [fs] - - if download_results: - fs_done = [f for f in fs if f.done] - fs_not_done = [f for f in fs if not f.done] - else: - fs_done = [f for f in fs if f.success or f.done] - fs_not_done = [f for f in fs if not (f.success or f.done)] + fs = _as_future_list(fs) + prefix = log_prefix(fs[0].executor_id) + fs_done, fs_not_done = _partition_futures(fs, download_results) if not fs_not_done: - logger.debug(f'ExecutorID {fs[0].executor_id} - All function activations are done') + logger.debug(f'{prefix} - All function activations are done') return fs_done, fs_not_done not_done_futures = fs_not_done if futures_from_executor_wait else fs - fs_to_wait = math.ceil(return_when * len(not_done_futures) / 100) - - if return_when == ALL_COMPLETED: - logger.info(f'ExecutorID {fs[0].executor_id} - Waiting for ' - f'{len(not_done_futures)} function activations to complete') - else: - txt = 'any' if return_when == ANY_COMPLETED else f'{return_when}%' - logger.info(f'ExecutorID {fs[0].executor_id} - Waiting for {txt} of ' - f'{len(not_done_futures)} function activations to complete') + _log_wait_start(prefix, return_when, len(not_done_futures)) if is_unix_system() and timeout is not None: - logger.debug(f'Setting waiting timeout to {timeout} seconds') - error_msg = 'Timeout of {timeout} seconds exceeded waiting for function activations to finish' - signal.signal(signal.SIGALRM, partial(timeout_handler, error_msg)) - signal.alarm(timeout) - - # Setup progress bar - pbar = None - if not is_lithops_worker() and show_progressbar and logger.getEffectiveLevel() != logging.DEBUG: - from tqdm.auto import tqdm - if not is_notebook(): - print() - pbar = tqdm(bar_format=' {l_bar}{bar}| {n_fmt}/{total_fmt} ', - total=fs_to_wait, disable=None) - pbar.update(min(len(fs_done), fs_to_wait)) + _set_wait_alarm(timeout) + + pbar = ( + _create_progressbar(fs_to_wait, len(fs_done)) + if show_progressbar else None + ) + started_monitors = [] + pool = None try: - executors_data = _create_executors_data_from_futures(fs, internal_storage) + executors_data = _create_executors_data_from_futures( + fs, internal_storage + ) if not job_monitor: - for executor_data in executors_data: - job_monitor = JobMonitor( - executor_id=executor_data.executor_id, - internal_storage=executor_data.internal_storage) - job_monitor.start(fs=executor_data.futures) - - sleep_sec = wait_dur_sec or WAIT_DUR_SEC if job_monitor.type == 'storage' \ - and job_monitor.storage_backend != 'localhost' else 0.1 + started_monitors = _start_job_monitors(executors_data) + # All of them run, but a single one drives the loop below and + # sets its poll interval. There is only one wait to pace + job_monitor = started_monitors[-1] + + sleep_sec = _poll_sleep_sec(job_monitor, wait_dur_sec) + pool = cf.ThreadPoolExecutor(max_workers=threadpool_size) + poll_kwargs = dict( + pbar=pbar, + throw_except=throw_except, + download_results=download_results, + threadpool_size=threadpool_size, + pool=pool, + ) if return_when == ALWAYS: for executor_data in executors_data: - _get_executor_data(fs, executor_data, pbar=pbar, - throw_except=throw_except, - download_results=download_results, - threadpool_size=threadpool_size) - else: - while not _check_done(fs, return_when, download_results): - if not job_monitor.is_alive(): - job_monitor.start(fs=fs) - for executor_data in executors_data: - new_data = _get_executor_data(fs, executor_data, pbar=pbar, - throw_except=throw_except, - download_results=download_results, - threadpool_size=threadpool_size) - time.sleep(0 if new_data else sleep_sec) - - except KeyboardInterrupt as e: - if download_results: - not_dones_call_ids = [(f.job_id, f.call_id) for f in fs if not f.done] + _get_executor_data(fs, executor_data, **poll_kwargs) else: - not_dones_call_ids = [(f.job_id, f.call_id) for f in fs if not f.success and not f.done] - msg = (f'Cancelled - Total Activations not done: {len(not_dones_call_ids)}') + _poll_until_done( + fs, executors_data, job_monitor, return_when, + download_results, sleep_sec, poll_kwargs + ) + + except KeyboardInterrupt: + _, not_done = _partition_futures(fs, download_results) if pbar: pbar.close() print() - logger.info(msg) - raise e - - except Exception as e: - raise e + logger.info(f'Cancelled - Total Activations not done: {len(not_done)}') + raise finally: + if pool is not None: + pool.shutdown(wait=True) + for monitor in started_monitors: + monitor.stop() if is_unix_system(): signal.alarm(0) if pbar and not pbar.disable: @@ -174,76 +289,85 @@ def wait(fs: Union[ResponseFuture, FuturesList, List[ResponseFuture]], if not is_notebook(): print() - if download_results: - fs_done = [f for f in fs if f.done] - fs_notdone = [f for f in fs if not f.done] - else: - fs_done = [f for f in fs if f.success or f.done] - fs_notdone = [f for f in fs if not f.success and not f.done] - - return fs_done, fs_notdone + return _partition_futures(fs, download_results) -def get_result(fs: Optional[Union[ResponseFuture, FuturesList, List[ResponseFuture]]] = None, - internal_storage: Optional[InternalStorage] = None, - throw_except: Optional[bool] = True, - timeout: Optional[int] = None, - threadpool_size: Optional[int] = THREADPOOL_SIZE, - wait_dur_sec: Optional[int] = None, - show_progressbar: Optional[bool] = True): +def get_result( + fs: Optional[ + Union[ResponseFuture, FuturesList, List[ResponseFuture]] + ] = None, + internal_storage: Optional[InternalStorage] = None, + throw_except: Optional[bool] = True, + timeout: Optional[int] = None, + threadpool_size: Optional[int] = THREADPOOL_SIZE, + wait_dur_sec: Optional[int] = None, + show_progressbar: Optional[bool] = True, +): """ For getting the results from all function activations :param fs: Futures list. Default None :param internal_storage: InternalStorage instance. Default None. - :param throw_except: Reraise exception if call raised. Default True. + :param throw_except: Reraise exception if call raised. + Default True. :param timeout: Timeout for waiting for results. - :param threadpool_size: Number of threads to use. Default 128 - :param wait_dur_sec: Time interval between each check. Default 1 second + :param threadpool_size: Number of threads to use. Default 64 + :param wait_dur_sec: Time interval between each check. + Default 1 second :param show_progressbar: whether or not to show the progress bar. :return: The result of the future/s """ - if type(fs) is not list and type(fs) is not FuturesList: - fs = [fs] + fs = _as_future_list(fs) + prefix = log_prefix(fs[0].executor_id) logger.info( - (f'ExecutorID {fs[0].executor_id} - Getting results from ' - f'{len(fs)} function activations') + f'{prefix} - Getting results from {len(fs)} function activations' ) - fs_done, _ = wait(fs=fs, throw_except=throw_except, - timeout=timeout, download_results=True, - internal_storage=internal_storage, - threadpool_size=threadpool_size, - wait_dur_sec=wait_dur_sec, - show_progressbar=show_progressbar) - result = [] - for f in [f for f in fs_done if not f.futures and f._produce_output]: - result.append(f.result(throw_except=throw_except)) + fs_done, _ = wait( + fs=fs, + throw_except=throw_except, + timeout=timeout, + download_results=True, + internal_storage=internal_storage, + threadpool_size=threadpool_size, + wait_dur_sec=wait_dur_sec, + show_progressbar=show_progressbar, + ) + result = [ + f.result(throw_except=throw_except) + for f in fs_done + if not f.futures and f._produce_output + ] - logger.debug(f"ExecutorID {fs[0].executor_id} - Finished getting results") + logger.debug(f'{prefix} - Finished getting results') return result def _create_executors_data_from_futures(fs, internal_storage): """ - Creates a dummy job necessary for the job monitor + Groups the futures by the executor that created them, and pairs every + group with the storage its statuses have to be read from """ + grouped = {} + for fut in fs: + grouped.setdefault(fut.executor_id, []).append(fut) + executor_jobs = [] - present_executors = {f.executor_id for f in fs} - - for executor_id in present_executors: - executor_data = SimpleNamespace() - executor_data.executor_id = executor_id - executor_data.futures = [f for f in fs if f.executor_id == executor_id] - f = executor_data.futures[0] - if internal_storage and internal_storage.backend == f._storage_config['backend']: + for executor_id, futures in grouped.items(): + executor_data = SimpleNamespace( + executor_id=executor_id, + futures=futures, + ) + backend = futures[0]._storage_config['backend'] + if internal_storage and internal_storage.backend == backend: executor_data.internal_storage = internal_storage else: - executor_data.internal_storage = InternalStorage(f._storage_config) - + executor_data.internal_storage = InternalStorage( + futures[0]._storage_config + ) executor_jobs.append(executor_data) return executor_jobs @@ -253,66 +377,81 @@ def _check_done(fs, return_when, download_results): """ Checks if return_when% of futures are ready or done """ - if download_results: - total_done = [f.done for f in fs].count(True) - else: - total_done = [f.success or f.done for f in fs].count(True) + total_done = sum( + 1 for f in fs if _future_is_complete(f, download_results) + ) if return_when == ANY_COMPLETED: return total_done >= 1 - else: - done_percentage = int(total_done * 100 / len(fs)) - return done_percentage >= return_when + + done_percentage = int(total_done * 100 / len(fs)) + return done_percentage >= return_when -def _get_executor_data(fs, exec_data, download_results, throw_except, threadpool_size, pbar): +def _ready_futures(exec_data, download_results): """ - Downloads all status/results from ready futures + Returns the futures of one executor that have data waiting on the other + side: their status has arrived, but the caller has not fetched it yet """ - if download_results: - callids_done = [(f.executor_id, f.job_id, f.call_id) for f in exec_data.futures if (f.ready or f.success)] - not_done_futures = [f for f in exec_data.futures if not f.done] + done_ids = { + _future_id(f) for f in exec_data.futures if f.ready or f.success + } + pending = (f for f in exec_data.futures if not f.done) else: - callids_done = [(f.executor_id, f.job_id, f.call_id) for f in exec_data.futures if f.ready] - not_done_futures = [f for f in exec_data.futures if not (f.success or f.done)] - - not_done_call_ids = set([(f.executor_id, f.job_id, f.call_id) for f in not_done_futures]) - new_callids_done = not_done_call_ids.intersection(callids_done) - - fs_to_wait_on = [] - for f in exec_data.futures: - if (f.executor_id, f.job_id, f.call_id) in new_callids_done: - fs_to_wait_on.append(f) + done_ids = {_future_id(f) for f in exec_data.futures if f.ready} + pending = (f for f in exec_data.futures if not (f.success or f.done)) + + ready_ids = {_future_id(f) for f in pending} & done_ids + return [f for f in exec_data.futures if _future_id(f) in ready_ids] + + +def _get_executor_data( + fs, + exec_data, + download_results, + throw_except, + threadpool_size, + pbar, + pool=None, +): + """ + Downloads the status, or the whole result, of every ready future of one + executor. Returns how many were fetched, so that the caller can tell a + productive poll from an empty one + """ + fs_to_wait_on = _ready_futures(exec_data, download_results) + if not fs_to_wait_on: + return 0 - def get_result(f): - f.result(throw_except=throw_except, internal_storage=exec_data.internal_storage) + storage = exec_data.internal_storage - def get_status(f): - f.status(throw_except=throw_except, internal_storage=exec_data.internal_storage) + def fetch(f): + if download_results: + f.result(throw_except=throw_except, internal_storage=storage) + else: + f.status(throw_except=throw_except, internal_storage=storage) - pool = cf.ThreadPoolExecutor(max_workers=threadpool_size) - if download_results: - list(pool.map(get_result, fs_to_wait_on)) + if pool is None: + with cf.ThreadPoolExecutor(max_workers=threadpool_size) as owned: + list(owned.map(fetch, fs_to_wait_on)) else: - list(pool.map(get_status, fs_to_wait_on)) - pool.shutdown() + list(pool.map(fetch, fs_to_wait_on)) if pbar: for f in fs_to_wait_on: - if (download_results and f.done) or \ - (not download_results and (f.success or f.done)): - if pbar.n < pbar.total: - pbar.update(1) + if _future_is_complete(f, download_results) and pbar.n < pbar.total: + pbar.update(1) pbar.refresh() - # Check for new futures - new_futures = list(chain(*[f._new_futures for f in fs_to_wait_on if f._new_futures])) + new_futures = list(chain.from_iterable( + f._new_futures for f in fs_to_wait_on if f._new_futures + )) if new_futures: fs.extend(new_futures) exec_data.futures.extend(new_futures) if pbar: - pbar.total = pbar.total + len(new_futures) + pbar.total += len(new_futures) pbar.refresh() return len(fs_to_wait_on) diff --git a/lithops/worker/handler.py b/lithops/worker/handler.py index 754d0cc42..39f11aa4f 100644 --- a/lithops/worker/handler.py +++ b/lithops/worker/handler.py @@ -17,12 +17,14 @@ import os import sys +import ast import zlib import time import json import uuid import base64 import pickle +import struct import logging import traceback import multiprocessing as mp @@ -30,26 +32,35 @@ from threading import Thread from tblib import pickling_support from types import SimpleNamespace -from multiprocessing.managers import SyncManager +from typing import Any, Callable, Dict, Optional, Tuple, Union from lithops.version import __version__ from lithops.config import extract_storage_config from lithops.storage import InternalStorage from lithops.worker.jobrunner import JobRunner -from lithops.worker.utils import LogStream, custom_redirection, \ - get_function_and_modules, get_function_data +from lithops.worker.utils import ( + LogStream, custom_redirection, get_function_and_modules, + get_function_data, SystemMonitor +) from lithops.constants import JOBS_PREFIX, LITHOPS_TEMP_DIR, MODULES_DIR -from lithops.utils import setup_lithops_logger, is_unix_system +from lithops.utils import ( + MONITORING_QUEUES_ENV, + setup_lithops_logger, + is_unix_system, +) from lithops.worker.status import create_call_status -from lithops.worker.utils import SystemMonitor pickling_support.install() logger = logging.getLogger(__name__) # Python 3.14 defaults to forkserver on Linux, which requires pickling Process -# arguments. Lithops relies on fork semantics for JobRunner subprocesses. -_MP_CTX = mp.get_context('fork') if is_unix_system() else None +# arguments. Lithops relies on fork semantics: both the JobRunner subprocess +# and the worker processes inherit the task from their parent. +_MP_CTX = mp.get_context('fork') if is_unix_system() else mp + +# A task, as passed from the work queue to a worker: (job, call_id, data) +Task = Tuple[SimpleNamespace, str, Any] class ShutdownSentinel: @@ -57,17 +68,151 @@ class ShutdownSentinel: pass -def create_job(payload: dict) -> SimpleNamespace: +class TaskJar: + """ + Work queue for forked worker processes, backed by a single POSIX pipe. + + Workers inherit the job and its data through fork, so the pipe only + carries a fixed size task index. Reading one is how a worker claims a + task: the kernel hands every token to exactly one reader, which balances + the load without a lock. Unlike multiprocessing.Queue and Manager this + needs no POSIX shared memory, so it also works on FaaS sandboxes that do + not provide /dev/shm. + """ + TOKEN = struct.Struct('!i') + # Writes up to PIPE_BUF are atomic, so a token is never split in two and + # readers stay aligned. 512 is the smallest PIPE_BUF POSIX allows. + MAX_ATOMIC_WRITE = 512 + + def __init__(self, job: SimpleNamespace): + self.job = job + # A worker rebinds job.data to the task it is running, so keep an + # independent reference to the full list of calls. + self.calls = list(zip(job.call_ids, job.data)) + self.read_fd, self.write_fd = os.pipe() + + def close_reader(self) -> None: + """ + Drops the parent's read end. Called once every worker is forked, so + that dispatch() fails instead of blocking if all the workers die. + """ + os.close(self.read_fd) + + def close_writer(self) -> None: + """ + Drops a worker's inherited write end. Called by every worker, as + otherwise the pipe never reaches EOF and no worker ever stops. + """ + os.close(self.write_fd) + + def dispatch(self) -> None: + """ + Offers every task to the workers, then closes the pipe so that they + see EOF and stop. Called by the parent process. + """ + tokens = b''.join(self.TOKEN.pack(i) for i in range(len(self.calls))) + view = memoryview(tokens) + try: + while view: + written = os.write(self.write_fd, view[:self.MAX_ATOMIC_WRITE]) + view = view[written:] + except BrokenPipeError: + logger.error('Worker processes exited before consuming all tasks') + finally: + self.close_writer() + + def get(self) -> Task: + """ + Claims the next task, blocking until one is available. Raises Empty + once the jar is exhausted. + """ + token = b'' + while len(token) < self.TOKEN.size: + chunk = os.read(self.read_fd, self.TOKEN.size - len(token)) + if not chunk: + raise Empty + token += chunk + + index, = self.TOKEN.unpack(token) + call_id, data = self.calls[index] + return self.job, call_id, data + + +def create_job(payload: Dict[str, Any]) -> SimpleNamespace: + """ + Builds a job out of an invocation payload, downloading the function, + the modules and the data it refers to + """ job = SimpleNamespace(**payload) storage_config = extract_storage_config(job.config) internal_storage = InternalStorage(storage_config) job.func = get_function_and_modules(job, internal_storage) job.data = get_function_data(job, internal_storage) - return job -def function_handler(payload): +def _fill_queue(job: SimpleNamespace, worker_processes: int) -> Queue: + """ + Loads every task of the job in a queue, followed by one sentinel per + worker. Every task is known upfront, so nothing is queued afterwards + """ + work_queue = Queue() + + for call_id, data in zip(job.call_ids, job.data): + work_queue.put((job, call_id, data)) + + for _ in range(worker_processes): + work_queue.put(ShutdownSentinel()) + + return work_queue + + +def _jar_worker(pid: int, jar: TaskJar) -> None: + """ + Entry point of a forked worker process + """ + jar.close_writer() + task_consumer(pid, jar) + + +def _run_process_pool(job: SimpleNamespace, worker_processes: int) -> None: + """ + Runs the tasks of the job in forked processes, each one claiming the next + task from the jar as soon as it is free + """ + jar = TaskJar(job) + workers = [] + + for pid in range(worker_processes): + worker = _MP_CTX.Process(target=_jar_worker, args=(pid, jar)) + workers.append(worker) + worker.start() + + jar.close_reader() + jar.dispatch() + + for worker in workers: + worker.join() + + +def _run_thread_pool(job: SimpleNamespace, worker_processes: int) -> None: + """ + Runs the tasks of the job in threads. Used where there is no fork, so + tasks share this interpreter instead of getting a process each + """ + work_queue = _fill_queue(job, worker_processes) + workers = [] + + for pid in range(worker_processes): + worker = Thread(target=task_consumer, args=(pid, work_queue)) + workers.append(worker) + worker.start() + + for worker in workers: + worker.join() + + +def function_handler(payload: Dict[str, Any]) -> None: """ Default function entry point called from Serverless backends """ @@ -75,37 +220,18 @@ def function_handler(payload): setup_lithops_logger(job.log_level) worker_processes = min(job.worker_processes, len(job.call_ids)) - logger.info(f'Tasks received: {len(job.call_ids)} - Worker processes: {worker_processes}') + logger.info( + f'Tasks received: {len(job.call_ids)} - ' + f'Worker processes: {worker_processes}' + ) if worker_processes == 1: - work_queue = Queue() - for call_id in job.call_ids: - data = job.data.pop(0) - work_queue.put((job, call_id, data)) - work_queue.put(ShutdownSentinel()) - python_queue_consumer(0, work_queue, ) + task_consumer(0, _fill_queue(job, worker_processes)) + elif is_unix_system(): + _run_process_pool(job, worker_processes) else: - manager = _MP_CTX.Manager() if _MP_CTX else SyncManager() - manager.start() - work_queue = manager.Queue() - job_runners = [] - - for call_id in job.call_ids: - data = job.data.pop(0) - work_queue.put((job, call_id, data)) + _run_thread_pool(job, worker_processes) - for pid in range(worker_processes): - work_queue.put(ShutdownSentinel()) - p = _MP_CTX.Process(target=python_queue_consumer, args=(pid, work_queue,)) - job_runners.append(p) - p.start() - - for runner in job_runners: - runner.join() - - manager.shutdown() - - # Delete modules path from syspath module_path = os.path.join(MODULES_DIR, job.job_key) if module_path in sys.path: sys.path.remove(module_path) @@ -113,17 +239,25 @@ def function_handler(payload): os.environ.pop('__LITHOPS_TOTAL_EXECUTORS', None) -def python_queue_consumer(pid, work_queue, initializer=None, callback=None): +def task_consumer( + pid: int, + work_queue: Union[Queue, TaskJar], + initializer: Optional[Callable] = None, + callback: Optional[Callable] = None +) -> None: """ - Listens to the job_queue and executes the individual job tasks + Runs tasks until the work queue is exhausted. + + Takes either a threading Queue, terminated by a ShutdownSentinel, or a + TaskJar, which raises Empty once its pipe reaches EOF. """ - logger.info(f'Worker process {pid} started') + logger.info(f'Worker {pid} started') + tasks_done = 0 + while True: try: - event = work_queue.get(block=True) - except Empty: - break - except BrokenPipeError: + event = work_queue.get() + except (Empty, BrokenPipeError): break if isinstance(event, ShutdownSentinel): @@ -133,16 +267,28 @@ def python_queue_consumer(pid, work_queue, initializer=None, callback=None): task.call_id = call_id task.data = data - initializer(pid, task) if initializer is not None else None + try: + if initializer: + initializer(pid, task) + + prepare_and_run_task(task) - prepare_and_run_task(task) + if callback: + callback(pid, task) + except Exception as e: + # Do not lose this worker for the tasks that are still pending + logger.error(f'Worker {pid} failed to run task {call_id}: {e}') - callback(pid, task) if callback is not None else None + tasks_done += 1 - logger.info(f'Worker process {pid} finished') + logger.info(f'Worker {pid} finished, {tasks_done} tasks executed') -def prepare_and_run_task(task): +def prepare_and_run_task(task: SimpleNamespace) -> None: + """ + Sets up the environment and the working directory of a single task, and + runs it with its output redirected to the task log + """ task.start_tstamp = time.time() if '__LITHOPS_ACTIVATION_ID' not in os.environ: @@ -155,24 +301,91 @@ def prepare_and_run_task(task): storage_backend = task.config['lithops']['storage'] bucket = task.config[storage_backend]['storage_bucket'] - task.task_dir = os.path.join(LITHOPS_TEMP_DIR, bucket, JOBS_PREFIX, task.job_key, task.call_id) + task.task_dir = os.path.join( + LITHOPS_TEMP_DIR, bucket, JOBS_PREFIX, task.job_key, task.call_id + ) task.log_file = os.path.join(task.task_dir, 'execution.log') task.stats_file = os.path.join(task.task_dir, 'job_stats.txt') os.makedirs(task.task_dir, exist_ok=True) - with open(task.log_file, 'a') as log_strem: - task.log_stream = LogStream(log_strem) + with open(task.log_file, 'a') as log_stream: + task.log_stream = LogStream(log_stream) with custom_redirection(task.log_stream): run_task(task) - # Unset specific job env vars for key in task.extra_env: os.environ.pop(key, None) -def run_task(task): +def _add_resource_usage(call_status, sys_monitor: SystemMonitor) -> None: + """ + Reports the CPU, network and memory that the task consumed + """ + cpu_info = sys_monitor.get_cpu_info() + call_status.add('worker_func_cpu_usage', cpu_info['usage']) + call_status.add('worker_func_cpu_system_time', round(cpu_info['system'], 8)) + call_status.add('worker_func_cpu_user_time', round(cpu_info['user'], 8)) + + net_io = sys_monitor.get_network_io() + call_status.add('worker_func_sent_net_io', net_io['sent']) + call_status.add('worker_func_recv_net_io', net_io['recv']) + + mem_info = sys_monitor.get_memory_info() + call_status.add('worker_func_rss', mem_info['rss']) + call_status.add('worker_func_vms', mem_info['vms']) + call_status.add('worker_func_uss', mem_info['uss']) + + +def _add_task_stats(call_status, stats_file: str) -> None: + """ + Reports the stats the JobRunner wrote, if it got as far as writing them + """ + if not os.path.exists(stats_file): + return + + with open(stats_file, 'r') as fid: + for line in fid.readlines(): + key, value = line.strip().split(" ", 1) + try: + call_status.add(key, float(value)) + except ValueError: + call_status.add(key, value) + if key in ['exception', 'exc_pickle_fail']: + call_status.add(key, ast.literal_eval(value)) + + +def _add_exception(call_status) -> None: """ - Runs a single job within a separate process + Prints the traceback to the task log and reports it back to the client. + Only valid while handling an exception + """ + print('----------------------- EXCEPTION !-----------------------') + traceback.print_exc(file=sys.stdout) + print('----------------------------------------------------------') + call_status.add('exception', True) + + pickled_exc = pickle.dumps(sys.exc_info()) + pickle.loads(pickled_exc) # fail here if the client could not unpickle it + call_status.add('exc_info', str(pickled_exc)) + + +def _add_logs(call_status, task: SimpleNamespace) -> None: + """ + Reports the task log, compressed, so that the client can replay it + """ + task.log_stream.flush() + if not os.path.isfile(task.log_file): + return + + with open(task.log_file, 'rb') as log_file: + compressed = zlib.compress(log_file.read()) + call_status.add('logs', base64.b64encode(compressed).decode()) + + +def run_task(task: SimpleNamespace) -> None: + """ + Runs a single task, with the user function isolated in a JobRunner + subprocess, and reports its status and its resource usage """ setup_lithops_logger(task.log_level) @@ -180,36 +393,50 @@ def run_task(task): logger.info(f"Lithops v{__version__} - Starting {backend} execution") logger.info(f"Execution ID: {task.job_key}/{task.call_id}") - env = task.extra_env - env['LITHOPS_CONFIG'] = json.dumps(task.config) - env['__LITHOPS_SESSION_ID'] = '-'.join([task.job_key, task.call_id]) - os.environ.update(env) + injected_env = { + 'LITHOPS_CONFIG': json.dumps(task.config), + '__LITHOPS_SESSION_ID': '-'.join([task.job_key, task.call_id]), + # An executor created by the user function reports to these queues as + # well as to its own, which is how a nested job reaches the client + MONITORING_QUEUES_ENV: json.dumps( + getattr(task, 'monitoring_queues', None) or [] + ), + } + os.environ.update(task.extra_env) + os.environ.update(injected_env) storage_config = extract_storage_config(task.config) internal_storage = InternalStorage(storage_config) call_status = create_call_status(task, internal_storage) - runtime_name = task.runtime_name - memory = task.runtime_memory - timeout = task.execution_timeout - if task.runtime_memory: - logger.debug(f'Runtime: {runtime_name} - Memory: {memory}MB - Timeout: {timeout} seconds') + logger.debug( + f'Runtime: {task.runtime_name} - Memory: {task.runtime_memory}MB - ' + f'Timeout: {task.execution_timeout} seconds' + ) else: - logger.debug(f'Runtime: {runtime_name} - Timeout: {timeout} seconds') + logger.debug( + f'Runtime: {task.runtime_name} - ' + f'Timeout: {task.execution_timeout} seconds' + ) - job_interruped = False + job_interrupted = False try: - # send init status event call_status.send_init_event() handler_conn, jobrunner_conn = _MP_CTX.Pipe() jobrunner = JobRunner(task, jobrunner_conn, internal_storage) logger.debug('Starting JobRunner process') - jrp = _MP_CTX.Process(target=jobrunner.run) if is_unix_system() else Thread(target=jobrunner.run) - - process_id = os.getpid() if is_unix_system() else mp.current_process().pid + jrp = ( + _MP_CTX.Process(target=jobrunner.run) + if is_unix_system() + else Thread(target=jobrunner.run) + ) + + process_id = ( + os.getpid() if is_unix_system() else mp.current_process().pid + ) sys_monitor = SystemMonitor(process_id) sys_monitor.start() @@ -219,77 +446,48 @@ def run_task(task): sys_monitor.stop() logger.debug('JobRunner process finished') - cpu_info = sys_monitor.get_cpu_info() - call_status.add('worker_func_cpu_usage', cpu_info['usage']) - call_status.add('worker_func_cpu_system_time', round(cpu_info['system'], 8)) - call_status.add('worker_func_cpu_user_time', round(cpu_info['user'], 8)) - - net_io = sys_monitor.get_network_io() - call_status.add('worker_func_sent_net_io', net_io['sent']) - call_status.add('worker_func_recv_net_io', net_io['recv']) - - mem_info = sys_monitor.get_memory_info() - call_status.add('worker_func_rss', mem_info['rss']) - call_status.add('worker_func_vms', mem_info['vms']) - call_status.add('worker_func_uss', mem_info['uss']) + _add_resource_usage(call_status, sys_monitor) if jrp.is_alive(): - # If process is still alive after jr.join(job_max_runtime), kill it try: jrp.terminate() except Exception: - # thread does not have terminate method + # Where there is no fork the JobRunner is a thread, which + # cannot be terminated. It is left behind on purpose pass - msg = ('Function exceeded maximum time of {} seconds and was ' - 'killed'.format(task.execution_timeout)) - raise TimeoutError('HANDLER', msg) + raise TimeoutError( + f'Function exceeded maximum time of {task.execution_timeout} ' + f'seconds and was killed' + ) if not handler_conn.poll(): - logger.error('No completion message received from JobRunner process') + # The JobRunner sends exactly one message when it finishes, so no + # message means it was killed. That is an OOM 99% of the times + logger.error( + 'No completion message received from JobRunner process' + ) logger.debug('Assuming memory overflow...') - # Only 1 message is returned by jobrunner when it finishes. - # If no message, this means that the jobrunner process was killed. - # 99% of times the jobrunner is killed due an OOM, so we assume here an OOM. - msg = 'Function exceeded maximum memory and was killed' - raise MemoryError('HANDLER', msg) - - if os.path.exists(task.stats_file): - with open(task.stats_file, 'r') as fid: - for line in fid.readlines(): - key, value = line.strip().split(" ", 1) - try: - call_status.add(key, float(value)) - except Exception: - call_status.add(key, value) - if key in ['exception', 'exc_pickle_fail']: - call_status.add(key, eval(value)) + raise MemoryError( + 'Function exceeded maximum memory and was killed' + ) + + _add_task_stats(call_status, task.stats_file) except KeyboardInterrupt: - job_interruped = True + job_interrupted = True logger.debug("Job interrupted") except Exception: - # internal runtime exceptions - print('----------------------- EXCEPTION !-----------------------') - traceback.print_exc(file=sys.stdout) - print('----------------------------------------------------------') - call_status.add('exception', True) - - pickled_exc = pickle.dumps(sys.exc_info()) - pickle.loads(pickled_exc) # this is just to make sure they can be unpickled - call_status.add('exc_info', str(pickled_exc)) + _add_exception(call_status) finally: - if not job_interruped: - call_status.add('worker_end_tstamp', time.time()) - - # Flush log stream and save it to the call status - task.log_stream.flush() - if os.path.isfile(task.log_file): - with open(task.log_file, 'rb') as lf: - log_str = base64.b64encode(zlib.compress(lf.read())).decode() - call_status.add('logs', log_str) + for key in injected_env: + os.environ.pop(key, None) + # An interrupted job is not reported: the client is gone anyway + if not job_interrupted: + call_status.add('worker_end_tstamp', time.time()) + _add_logs(call_status, task) call_status.send_finish_event() logger.info("Finished") diff --git a/lithops/worker/invoker.py b/lithops/worker/invoker.py index 023277dcc..d4e7813c9 100644 --- a/lithops/worker/invoker.py +++ b/lithops/worker/invoker.py @@ -14,54 +14,59 @@ # limitations under the License. # import os +import json import time import logging from types import SimpleNamespace +from typing import Any, Dict from lithops.serverless import ServerlessHandler from lithops.monitor import JobMonitor from lithops.storage import InternalStorage from lithops.config import extract_serverless_config, extract_storage_config from lithops.invokers import FaaSInvoker +from lithops.utils import MONITORING_QUEUES_ENV, monitoring_queues logger = logging.getLogger(__name__) -def function_invoker(job_payload): +def function_invoker(job_payload: Dict[str, Any]) -> None: """ - Method used as a remote invoker + Entry point of the remote invoker: invokes a whole job from a worker, + instead of from the client """ config = job_payload['config'] job = SimpleNamespace(**job_payload['job']) - env = {'LITHOPS_WORKER': 'True', 'PYTHONUNBUFFERED': 'True', - '__LITHOPS_SESSION_ID': job.job_key} - os.environ.update(env) + os.environ.update({ + 'LITHOPS_WORKER': 'True', + 'PYTHONUNBUFFERED': 'True', + '__LITHOPS_SESSION_ID': job.job_key, + # The job this invoker spawns reports to the queues of the client, and + # an executor created here extends that chain rather than replacing it + MONITORING_QUEUES_ENV: json.dumps( + monitoring_queues(job.executor_id) + ), + }) backend = config['lithops']['backend'] config[backend]['invoke_pool_threads'] = 128 - # Create the internal_storage handler storage_config = extract_storage_config(config) internal_storage = InternalStorage(storage_config) - # Create the compute handler serverless_config = extract_serverless_config(config) compute_handler = ServerlessHandler(serverless_config, storage_config) - # Create the monitoring system monitoring_backend = config['lithops']['monitoring'].lower() - monitoring_config = config.get(monitoring_backend) - job_monitor = JobMonitor( executor_id=job.executor_id, internal_storage=internal_storage, backend=monitoring_backend, - config=monitoring_config + config=config.get(monitoring_backend) ) - # Create the invoker invoker = FaaSRemoteInvoker( config, job.executor_id, @@ -74,12 +79,13 @@ def function_invoker(job_payload): class FaaSRemoteInvoker(FaaSInvoker): """ - Module responsible to perform the invocations against the serverless compute backend + Module responsible to perform the invocations against the serverless + compute backend """ - def run_job(self, job): + def run_job(self, job: SimpleNamespace) -> None: """ - Run a job + Invokes every task of the job and waits until they are all submitted """ futures = self._run_job(job) self.job_monitor.start( @@ -89,11 +95,14 @@ def run_job(self, job): generate_tokens=True ) + # stop() drops whatever is still pending, so wait until the async + # invokers have picked every chunk up before stopping them while self.pending_calls_q.qsize() > 0: time.sleep(1) - self.job_monitor.stop() # Stop job monitor thread - self.stop() # Stop async invokers threads - time.sleep(5) + self.job_monitor.stop() + # Waits for the invocations still in flight, which this worker must not + # be frozen in the middle of + self.stop(wait=True) logger.info('Remote Invoker Finished') diff --git a/lithops/worker/jobrunner.py b/lithops/worker/jobrunner.py index 44a11ae6d..6d893822a 100644 --- a/lithops/worker/jobrunner.py +++ b/lithops/worker/jobrunner.py @@ -27,9 +27,12 @@ import requests import traceback from pydoc import locate +from types import SimpleNamespace +from typing import Any, Callable, Dict, Optional, Tuple from lithops.worker.utils import peak_memory +# Importing numpy here makes numpy types pickle-compatible in the worker. try: import numpy as np np.__version__ @@ -39,172 +42,337 @@ from lithops.storage import Storage from lithops.wait import wait from lithops.future import ResponseFuture -from lithops.utils import WrappedStreamingBody, sizeof_fmt, \ - is_object_processing_function, FuturesList, verify_args -from lithops.utils import WrappedStreamingBodyPartition +from lithops.utils import ( + WrappedStreamingBody, sizeof_fmt, is_object_processing_function, + FuturesList, verify_args, WrappedStreamingBodyPartition +) from lithops.util.metrics import PrometheusExporter from lithops.storage.utils import create_output_key logger = logging.getLogger(__name__) +# Results below this size are written to the stats file, which travels back +# with the call status, instead of being uploaded to storage on their own +_MAX_INLINE_RESULT_SIZE = 8 * 1024 + + +def _prepost(func): + """Runs the PRE_RUN / POST_RUN callables from the environment around func""" + def call(env_var): + if env_var in os.environ: + method = locate(os.environ[env_var]) + method() + + def wrapper_decorator(*args, **kwargs): + call('PRE_RUN') + value = func(*args, **kwargs) + call('POST_RUN') + return value + return wrapper_decorator + class JobStats: + """ + Line based stats file, written as the task progresses and read back by + the handler once the JobRunner is done + """ - def __init__(self, stats_filename): + def __init__(self, stats_filename: str): self.stats_filename = stats_filename self.stats_fid = open(stats_filename, 'w') - def write(self, key, value): - self.stats_fid.write("{} {}\n".format(key, value)) + def write(self, key: str, value: Any) -> None: + """Appends one stat, flushed so that it survives a killed worker""" + self.stats_fid.write(f"{key} {value}\n") self.stats_fid.flush() + def close(self) -> None: + """Closes the stats file, unless it is closed already""" + if getattr(self, 'stats_fid', None) and not self.stats_fid.closed: + self.stats_fid.close() + def __del__(self): - self.stats_fid.close() + self.close() + + +def _get_function_name(func: Callable) -> str: + """Returns the name of a function, or of the class of a callable object""" + if inspect.isfunction(func) or inspect.ismethod(func): + return func.__name__ + return type(func).__name__ + + +def _returns_futures(result: Any) -> bool: + """ + Tells whether the function returned futures to chain, instead of data. + A list is only inspected by its first element, as the client does + """ + if isinstance(result, (ResponseFuture, FuturesList)): + return True + return ( + isinstance(result, list) + and len(result) > 0 + and isinstance(result[0], ResponseFuture) + ) class JobRunner: + """ + Runs the user function of a single task, isolated from the handler, and + reports the result and the stats through the task stats file + """ - def __init__(self, job, jobrunner_conn, internal_storage): + def __init__(self, job: SimpleNamespace, jobrunner_conn, internal_storage): self.job = job self.jobrunner_conn = jobrunner_conn self.internal_storage = internal_storage self.lithops_config = job.config - self.output_key = create_output_key(job.executor_id, job.job_id, job.call_id) - - # Setup stats class + self.output_key = create_output_key( + job.executor_id, job.job_id, job.call_id + ) self.stats = JobStats(self.job.stats_file) - # Setup prometheus for live metrics prom_enabled = self.lithops_config['lithops'].get('telemetry') prom_config = self.lithops_config.get('prometheus', {}) self.prometheus = PrometheusExporter(prom_enabled, prom_config) - def _fill_optional_args(self, function, data): + def _prom_labels( + self, fn_name: Optional[str] + ) -> Tuple[Tuple[str, str], ...]: + return ( + ('job_id', self.job.job_key), + ('call_id', '-'.join([self.job.job_key, self.job.call_id])), + ('function_name', fn_name or 'undefined') + ) + + def _create_ibm_cos_client(self): + """Creates the boto3 client injected as the ibm_cos parameter""" + if 'ibm_cos' not in self.lithops_config: + raise Exception( + 'Cannot create the ibm_cos client: missing configuration' + ) + + if self.internal_storage.backend == 'ibm_cos': + return self.internal_storage.get_client() + + return Storage( + config=self.lithops_config, backend='ibm_cos' + ).get_client() + + def _create_rabbitmq_connection(self): + """Creates the connection injected as the rabbitmq parameter""" + if 'rabbitmq' not in self.lithops_config: + raise Exception( + 'Cannot create the rabbitmq client: missing configuration' + ) + + rabbit_amqp_url = self.lithops_config['rabbitmq'].get('amqp_url') + return pika.BlockingConnection(pika.URLParameters(rabbit_amqp_url)) + + def _fill_optional_args( + self, function: Callable, data: Dict[str, Any] + ) -> None: """ - Fills in those reserved, optional parameters that might be write to the function signature + Fills in those reserved, optional parameters that might be written to + the function signature """ func_sig = inspect.signature(function) if len(data) == 1 and 'future' in data: # Function chaining feature - out = [data.pop('future').result(internal_storage=self.internal_storage)] + out = [ + data.pop('future').result( + internal_storage=self.internal_storage + ) + ] data.update(verify_args(function, out, None)[0]) if 'ibm_cos' in func_sig.parameters: - if 'ibm_cos' in self.lithops_config: - if self.internal_storage.backend == 'ibm_cos': - ibm_boto3_client = self.internal_storage.get_client() - else: - ibm_boto3_client = Storage(config=self.lithops_config, backend='ibm_cos').get_client() - data['ibm_cos'] = ibm_boto3_client - else: - raise Exception('Cannot create the ibm_cos client: missing configuration') + data['ibm_cos'] = self._create_ibm_cos_client() if 'storage' in func_sig.parameters: data['storage'] = self.internal_storage.storage if 'rabbitmq' in func_sig.parameters: - if 'rabbitmq' in self.lithops_config: - rabbit_amqp_url = self.lithops_config['rabbitmq'].get('amqp_url') - params = pika.URLParameters(rabbit_amqp_url) - connection = pika.BlockingConnection(params) - data['rabbitmq'] = connection - else: - raise Exception('Cannot create the rabbitmq client: missing configuration') + data['rabbitmq'] = self._create_rabbitmq_connection() if 'id' in func_sig.parameters: data['id'] = int(self.job.call_id) - def _wait_futures(self, data): + def _wait_futures(self, data: Dict[str, Any]) -> None: + """ + Replaces the futures a reduce function receives by their results, + blocking until every one of them is done + """ logger.info('Reduce function: waiting for map results') - fut_list = list(data.values())[0] + key = next(iter(data)) + fut_list = data[key] wait(fut_list, self.internal_storage, download_results=True) results = [f.result() for f in fut_list if f.done and not f.futures] fut_list.clear() - data[next(iter(data))] = results - - def _load_object(self, data): - """ - Loads the object in case of object processing - """ - extra_get_args = {} - obj = data['obj'] + data[key] = results + def _open_object_stream(self, obj: Any, extra_get_args: Dict[str, Any]): + """Opens the object to process, wherever it lives""" if hasattr(obj, 'bucket') and not hasattr(obj, 'path'): - logger.info(f'Getting dataset from {obj.backend}://{obj.bucket}/{obj.key}') + logger.info( + f'Getting dataset from {obj.backend}://{obj.bucket}/{obj.key}' + ) if obj.backend == self.internal_storage.backend: storage = self.internal_storage.storage else: - storage = Storage(config=self.lithops_config, backend=obj.backend) - if obj.data_byte_range is not None: - extra_get_args['Range'] = 'bytes={}-{}'.format(*obj.data_byte_range) - stream = storage.get_object(obj.bucket, obj.key, stream=True, extra_get_args=extra_get_args) - stream_body = stream + storage = Storage( + config=self.lithops_config, backend=obj.backend + ) + return storage.get_object( + obj.bucket, obj.key, stream=True, extra_get_args=extra_get_args + ) - elif hasattr(obj, 'url'): + if hasattr(obj, 'url'): logger.info(f'Getting dataset from {obj.url}') - if obj.data_byte_range is not None: - extra_get_args['Range'] = 'bytes={}-{}'.format(*obj.data_byte_range) - stream = requests.get(obj.url, headers=extra_get_args, stream=True).raw - stream_body = stream - - elif hasattr(obj, 'path'): - logger.info(f'Getting dataset from {obj.path}') - with open(obj.path, "rb") as f: - if obj.data_byte_range is not None: - first_byte, last_byte = obj.data_byte_range - f.seek(first_byte) - stream = io.BytesIO(f.read(last_byte - first_byte + 1)) - else: - stream = io.BytesIO(f.read()) - stream_body = stream + return requests.get( + obj.url, headers=extra_get_args, stream=True + ).raw + + logger.info(f'Getting dataset from {obj.path}') + with open(obj.path, "rb") as f: + if obj.data_byte_range is None: + return io.BytesIO(f.read()) + first_byte, last_byte = obj.data_byte_range + f.seek(first_byte) + return io.BytesIO(f.read(last_byte - first_byte + 1)) + def _load_object(self, data: Dict[str, Any]) -> None: + """ + Opens the object to process as a stream, and narrows its byte range + down to the chunk that this task is responsible for + """ + obj = data['obj'] + extra_get_args = {} if obj.data_byte_range is not None: - if obj.newline is None: - stream_body = WrappedStreamingBody(stream, obj.chunk_size) - else: - stream_body = WrappedStreamingBodyPartition(stream, obj.chunk_size, obj.data_byte_range, obj.newline) + first_byte, last_byte = obj.data_byte_range + extra_get_args['Range'] = f'bytes={first_byte}-{last_byte}' - obj.data_stream = stream_body + stream = self._open_object_stream(obj, extra_get_args) - if obj.data_byte_range is not None: - first_byte, last_byte = obj.data_byte_range + if obj.data_byte_range is None: + obj.data_stream = stream + first_byte = 0 + last_byte = obj.chunk_size - 1 + obj.data_byte_range = (first_byte, last_byte) + else: + if obj.newline is None: + obj.data_stream = WrappedStreamingBody(stream, obj.chunk_size) + else: + obj.data_stream = WrappedStreamingBodyPartition( + stream, obj.chunk_size, obj.data_byte_range, obj.newline + ) if last_byte - first_byte > obj.chunk_size: last_byte = first_byte + obj.chunk_size - 1 obj.data_byte_range = (first_byte, last_byte) - else: - first_byte = 0 - last_byte = obj.chunk_size - 1 - obj.data_byte_range = (0, last_byte) - - logger.info(f'Chunk: {obj.part}/{obj.total_parts} - Size: {obj.chunk_size} - Range: {first_byte}-{last_byte}') - - # Decorator to execute pre-run and post-run functions provided via environment variables - def prepost(func): - def call(envVar): - if envVar in os.environ: - method = locate(os.environ[envVar]) - method() - - def wrapper_decorator(*args, **kwargs): - call('PRE_RUN') - value = func(*args, **kwargs) - call('POST_RUN') - return value - return wrapper_decorator - - @prepost - def run(self): + + logger.info( + f'Chunk: {obj.part}/{obj.total_parts} - Size: {obj.chunk_size} - ' + f'Range: {first_byte}-{last_byte}' + ) + + def _write_function_stats( + self, start_tstamp: float, end_tstamp: float + ) -> None: + """ + Reports how long the user function took, with a result size that + _write_result overwrites if the function returned anything + """ + self.stats.write('worker_func_start_tstamp', start_tstamp) + self.stats.write('worker_func_end_tstamp', end_tstamp) + self.stats.write( + 'worker_func_exec_time', round(end_tstamp - start_tstamp, 8) + ) + self.stats.write('func_result_size', 0) + + def _write_result(self, result: Any) -> Optional[bytes]: """ - Runs the function + Reports the result of the function, and returns the pickled result + back when it is too big to travel with the call status + """ + if result is None: + return None + + if _returns_futures(result): + self.stats.write('new_futures', pickle.dumps(result)) + return None + + logger.debug("Pickling result") + pickled_output = pickle.dumps(result) + self.stats.write('func_result_size', len(pickled_output)) + + if len(pickled_output) >= _MAX_INLINE_RESULT_SIZE: + return pickled_output + + self.stats.write('result', pickled_output) + self.stats.write("worker_result_upload_time", 0) + return None + + def _write_exception(self) -> None: + """ + Prints the traceback to the task log and reports the exception, so + that the client can re-raise it. Only valid while handling one + """ + self.stats.write("exception", True) + exc_type, exc_value, exc_traceback = sys.exc_info() + print('----------------------- EXCEPTION !-----------------------') + traceback.print_exc(file=sys.stdout) + print('----------------------------------------------------------') + + try: + logger.debug("Pickling exception") + pickled_exc = pickle.dumps((exc_type, exc_value, exc_traceback)) + pickle.loads(pickled_exc) + + except Exception as pickle_exception: + # Shockingly often, modules like subprocess don't properly call + # the base Exception.__init__, which results in them being + # unpickleable. Report the pieces that do pickle instead of + # losing the exception altogether + self.stats.write("exc_pickle_fail", True) + pickled_exc = pickle.dumps({ + 'exc_type': str(exc_type), + 'exc_value': str(exc_value), + 'exc_traceback': exc_traceback, + 'pickle_exception': pickle_exception, + }) + pickle.loads(pickled_exc) + + self.stats.write("exc_info", str(pickled_exc)) + + def _upload_result(self, pickled_output: bytes) -> None: + """ + Uploads a result too big to travel with the call status, and reports + how long the upload took + """ + upload_start_tstamp = time.time() + logger.info( + f"Storing function result - " + f"Size: {sizeof_fmt(len(pickled_output))}" + ) + self.internal_storage.put_data(self.output_key, pickled_output) + upload_end_tstamp = time.time() + self.stats.write( + "worker_result_upload_time", + round(upload_end_tstamp - upload_start_tstamp, 8) + ) + + @_prepost + def run(self) -> None: + """ + Runs the user function and reports everything the client needs: its + result or its exception, its stats and its peak memory """ - # self.stats.write('worker_jobrunner_start_tstamp', time.time()) self.stats.write('worker_peak_memory_start', peak_memory()) logger.debug("Process started") - result = None - exception = False fn_name = None + pending_output = None try: func = pickle.loads(self.job.func) @@ -217,21 +385,15 @@ def run(self): self._fill_optional_args(func, data) - fn_name = func.__name__ if inspect.isfunction(func) \ - or inspect.ismethod(func) else type(func).__name__ - + fn_name = _get_function_name(func) self.prometheus.send_metric( name='function_start', value=time.time(), type='gauge', - labels=( - ('job_id', self.job.job_key), - ('call_id', '-'.join([self.job.job_key, self.job.call_id])), - ('function_name', fn_name or 'undefined') - ) + labels=self._prom_labels(fn_name) ) - logger.info(f"Going to execute '{str(fn_name)}()'") + logger.info(f"Going to execute '{fn_name}()'") print('---------------------- FUNCTION LOG ----------------------') function_start_tstamp = time.time() args, kwargs = _prepare_args(func, data) @@ -240,82 +402,38 @@ def run(self): print('----------------------------------------------------------') logger.info("Success function execution") - self.stats.write('worker_func_start_tstamp', function_start_tstamp) - self.stats.write('worker_func_end_tstamp', function_end_tstamp) - self.stats.write('worker_func_exec_time', round(function_end_tstamp - function_start_tstamp, 8)) - self.stats.write('func_result_size', 0) - - if result is not None: - # Check for new futures - if isinstance(result, ResponseFuture) or isinstance(result, FuturesList) \ - or (type(result) is list and len(result) > 0 and isinstance(result[0], ResponseFuture)): - self.stats.write('new_futures', pickle.dumps(result)) - result = None - else: - logger.debug("Pickling result") - pickled_output = pickle.dumps(result) - pickled_output_size = len(pickled_output) - self.stats.write('func_result_size', pickled_output_size) - if pickled_output_size < 8 * 1024: # 8KB - self.stats.write('result', pickled_output) - self.stats.write("worker_result_upload_time", 0) - result = None + self._write_function_stats( + function_start_tstamp, function_end_tstamp + ) + pending_output = self._write_result(result) except Exception: - exception = True - self.stats.write("exception", True) - exc_type, exc_value, exc_traceback = sys.exc_info() - print('----------------------- EXCEPTION !-----------------------') - traceback.print_exc(file=sys.stdout) - print('----------------------------------------------------------') - - try: - logger.debug("Pickling exception") - pickled_exc = pickle.dumps((exc_type, exc_value, exc_traceback)) - pickle.loads(pickled_exc) # this is just to make sure they can be unpickled - self.stats.write("exc_info", str(pickled_exc)) - - except Exception as pickle_exception: - # Shockingly often, modules like subprocess don't properly - # call the base Exception.__init__, which results in them - # being unpickleable. As a result, we actually wrap this in a try/catch block - # and more-carefully handle the exceptions if any part of this save / test-reload - # fails - self.stats.write("exc_pickle_fail", True) - pickled_exc = pickle.dumps({'exc_type': str(exc_type), - 'exc_value': str(exc_value), - 'exc_traceback': exc_traceback, - 'pickle_exception': pickle_exception}) - pickle.loads(pickled_exc) # this is just to make sure it can be unpickled - self.stats.write("exc_info", str(pickled_exc)) + self._write_exception() finally: - # self.stats.write('worker_jobrunner_end_tstamp', time.time()) self.stats.write('worker_peak_memory_end', peak_memory()) self.prometheus.send_metric( name='function_end', value=time.time(), type='gauge', - labels=( - ('job_id', self.job.job_key), - ('call_id', '-'.join([self.job.job_key, self.job.call_id])), - ('function_name', fn_name or 'undefined') - ) + labels=self._prom_labels(fn_name) ) - if result is not None and not exception: - output_upload_start_tstamp = time.time() - logger.info(f"Storing function result - Size: {sizeof_fmt(len(pickled_output))}") - self.internal_storage.put_data(self.output_key, pickled_output) - output_upload_end_tstamp = time.time() - self.stats.write("worker_result_upload_time", round(output_upload_end_tstamp - output_upload_start_tstamp, 8)) + if pending_output is not None: + self._upload_result(pending_output) + self.jobrunner_conn.send("Finished") logger.info("Process finished") + self.stats.close() -def _prepare_args(func, data): - # Convert the "data" envelope into normal *args/**kwargs, - # respecting the actual var-length parameter names of `func`. +def _prepare_args( + func: Callable, data: Dict[str, Any] +) -> Tuple[Any, Dict[str, Any]]: + """ + Converts the data envelope into normal args and kwargs, respecting the + actual var-length parameter names of func + """ func_sig = inspect.signature(func) var_pos_name = None var_kw_name = None @@ -328,11 +446,20 @@ def _prepare_args(func, data): payload = dict(data) - # Extract var-positional argument value if present - args = payload.pop(var_pos_name) or () if var_pos_name in payload else () - # Extract var-keyword argument value if present - kwargs = payload.pop(var_kw_name) or {} if var_kw_name in payload else {} - # Any remaining keys become normal keyword arguments + if var_pos_name is not None and var_pos_name in payload: + args = payload.pop(var_pos_name) + if args is None: + args = () + else: + args = () + + if var_kw_name is not None and var_kw_name in payload: + kwargs = payload.pop(var_kw_name) + if kwargs is None: + kwargs = {} + else: + kwargs = {} + kwargs.update(payload) return args, kwargs diff --git a/lithops/worker/status.py b/lithops/worker/status.py index b6880782f..7d7a2de27 100644 --- a/lithops/worker/status.py +++ b/lithops/worker/status.py @@ -1,16 +1,29 @@ +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + import os -import ast import pika import json import time import logging from tblib import pickling_support from contextlib import contextmanager +from types import SimpleNamespace +from typing import Any, Iterator -import lithops.worker -from lithops.utils import sizeof_fmt -from lithops.storage.utils import create_status_key, \ - create_init_key +from lithops.utils import monitoring_queue_name, sizeof_fmt +from lithops.storage.utils import create_status_key, create_init_key pickling_support.install() @@ -18,17 +31,25 @@ logger = logging.getLogger(__name__) -def create_call_status(job, internal_storage): - """ Creates a call status class based on the monitoring backend""" - monitoring_backend = job.config['lithops']['monitoring'] - Status = getattr(lithops.worker.status, '{}CallStatus' - .format(monitoring_backend.capitalize())) - return Status(job, internal_storage) +def create_call_status(job: SimpleNamespace, internal_storage) -> 'CallStatus': + """Creates a call status class based on the monitoring backend""" + monitoring_backend = job.config['lithops']['monitoring'].lower() + try: + status_cls = _STATUS_CLASSES[monitoring_backend] + except KeyError as exc: + raise ValueError( + f'Unknown monitoring backend: {monitoring_backend}' + ) from exc + return status_cls(job, internal_storage) class CallStatus: + """ + Status of a single call, reported to the client both when the task starts + and when it finishes + """ - def __init__(self, job, internal_storage): + def __init__(self, job: SimpleNamespace, internal_storage): self.job = job self.config = job.config self.internal_storage = internal_storage @@ -45,32 +66,34 @@ def __init__(self, job, internal_storage): 'chunksize': job.chunksize } - if ast.literal_eval(os.environ.get('WARM_CONTAINER', 'False')): - self.status['worker_cold_start'] = False - else: - self.status['worker_cold_start'] = True + is_warm = os.environ.get('WARM_CONTAINER', '').lower() in { + '1', 'true', 'yes' + } + self.status['worker_cold_start'] = not is_warm + if not is_warm: os.environ['WARM_CONTAINER'] = 'True' - def add(self, key, value): + def add(self, key: str, value: Any) -> None: """ Adds data to the call status""" self.status[key] = value - def send_init_event(self): + def send_init_event(self) -> None: """ Sends the init event""" self.status['type'] = '__init__' self._send() - def send_finish_event(self): + def send_finish_event(self) -> None: """ Sends the finish event""" self.status['type'] = '__end__' self._send() class StorageCallStatus(CallStatus): + """Reports the status of a call by writing it to the Object Storage""" - def _send(self): + def _send(self) -> None: """ - Send the status event to the Object Storage + Sends the status event to the Object Storage """ executor_id = self.status['executor_id'] job_id = self.status['job_id'] @@ -84,58 +107,94 @@ def _send(self): elif self.status['type'] == '__end__': status_key = create_status_key(executor_id, job_id, call_id) dmpd_response_status = json.dumps(self.status) - drs = sizeof_fmt(len(dmpd_response_status)) - logger.info("Storing execution stats - Size: {}".format(drs)) + logger.info( + f"Storing execution stats - " + f"Size: {sizeof_fmt(len(dmpd_response_status))}" + ) self.internal_storage.put_data(status_key, dmpd_response_status) class RabbitmqCallStatus(StorageCallStatus): + """ + Reports the status of a call by publishing it to RabbitMQ, which reaches + the client faster, and falls back to the Object Storage at the end + """ + MAX_ATTEMPTS = 5 - def __init__(self, job, internal_storage): + def __init__(self, job: SimpleNamespace, internal_storage): super().__init__(job, internal_storage) rabbit_amqp_url = self.config['rabbitmq'].get('amqp_url') self.pikaparams = pika.URLParameters(rabbit_amqp_url) @contextmanager - def _create_channel(self): + def _create_channel(self) -> Iterator[Any]: """ - Creates a rabbitmq channel + Creates a rabbitmq channel, closed along with its connection """ - self.connection = pika.BlockingConnection(self.pikaparams) - self.channel = self.connection.channel() + connection = pika.BlockingConnection(self.pikaparams) + channel = connection.channel() try: - yield self.channel + yield channel finally: - self.channel.close() - self.connection.close() + channel.close() + connection.close() - def _send(self): + def _queue_names(self): """ - Send the status event to RabbitMQ + Returns the name of every queue this status has to be published to, + which the client worked out and sent along with the job. + + The fallback only reaches the queue of this very executor: a payload + without the chain cannot say which executors are waiting further up, + so a nested job would go unnoticed by its ancestors """ - dmpd_response_status = json.dumps(self.status) - drs = sizeof_fmt(len(dmpd_response_status)) + queues = getattr(self.job, 'monitoring_queues', None) + if queues: + return list(queues) - status_sent = False - output_query_count = 0 + logger.warning( + 'The job carries no monitoring queues, reporting only to the ' + f'queue of {self.job.executor_id}' + ) + return [monitoring_queue_name(self.job.executor_id)] - queues = [] - executor_keys = self.job.executor_id.split('-') - for k in range(int(len(executor_keys) / 2)): - qname = 'lithops-{}'.format('-'.join(executor_keys[0:k * 3 + 2])) - queues.append(qname) + def _send(self) -> None: + """ + Sends the status event to RabbitMQ + """ + dmpd_response_status = json.dumps(self.status) + queues = self._queue_names() + exc = None - while not status_sent and output_query_count < 5: - output_query_count = output_query_count + 1 + for _ in range(self.MAX_ATTEMPTS): try: - with self._create_channel() as ch: + with self._create_channel() as channel: for queue in queues: - ch.basic_publish(exchange='', routing_key=queue, body=dmpd_response_status) - logger.info("Execution status sent to RabbitMQ - Size: {}".format(drs)) - status_sent = True - except Exception: + channel.basic_publish( + exchange='', + routing_key=queue, + body=dmpd_response_status + ) + logger.info( + f"Execution status sent to RabbitMQ - " + f"Size: {sizeof_fmt(len(dmpd_response_status))}" + ) + break + except Exception as e: + exc = e time.sleep(0.2) + else: + logger.error( + f"Could not send the execution status to RabbitMQ after " + f"{self.MAX_ATTEMPTS} attempts: {exc}" + ) if self.status['type'] == '__end__': super()._send() + + +_STATUS_CLASSES = { + 'storage': StorageCallStatus, + 'rabbitmq': RabbitmqCallStatus, +} diff --git a/lithops/worker/utils.py b/lithops/worker/utils.py index fb20c1d37..56751de63 100644 --- a/lithops/worker/utils.py +++ b/lithops/worker/utils.py @@ -15,17 +15,22 @@ # import os +import posixpath import sys +import ast import pkgutil import logging import pickle import platform import subprocess from contextlib import contextmanager +from types import SimpleNamespace +from typing import Any, Dict, List, Optional, Union from lithops.version import __version__ as lithops_ver -from lithops.utils import sizeof_fmt, is_unix_system, b64str_to_bytes -from lithops.constants import MODULES_DIR, SA_INSTALL_DIR, LITHOPS_TEMP_DIR +from lithops.utils import sizeof_fmt, is_unix_system +from lithops.constants import MODULES_DIR, SA_INSTALL_DIR +from lithops.job.serialize import write_module_data try: import psutil @@ -43,19 +48,21 @@ import ps_mem -def get_function_and_modules(job, internal_storage): +def get_function_and_modules(job: SimpleNamespace, internal_storage) -> bytes: """ - Gets the function and modules from storage + Gets the pickled function from storage, and writes the modules it depends + on where the interpreter can import them """ logger.info("Getting function and modules") backend = job.config['lithops']['backend'] - func_path = '/'.join([LITHOPS_TEMP_DIR, job.func_key]) - func_obj = None if job.config[backend].get('runtime_include_function'): - logger.info("Runtime include function feature activated. Loading " - "function/mods from local runtime") - func_path = '/'.join([SA_INSTALL_DIR, job.func_key]) + logger.info( + "Runtime include function feature activated. Loading " + "function/mods from local runtime" + ) + # Custom runtimes live on Linux images under /opt/lithops. + func_path = posixpath.join(SA_INSTALL_DIR, job.func_key) with open(func_path, "rb") as f: func_obj = f.read() else: @@ -69,47 +76,37 @@ def get_function_and_modules(job, internal_storage): logger.info(f"Writing function dependencies to {module_path}") os.makedirs(module_path, exist_ok=True) sys.path.append(module_path) - - for m_filename, m_data in loaded_func_all['module_data'].items(): - m_path = os.path.dirname(m_filename) - - if len(m_path) > 0 and m_path[0] == "/": - m_path = m_path[1:] - to_make = os.path.join(module_path, m_path) - try: - os.makedirs(to_make) - except OSError as e: - if e.errno == 17: - pass - else: - raise e - full_filename = os.path.join(to_make, os.path.basename(m_filename)) - # logger.debug('Writing {}'.format(full_filename)) - - with open(full_filename, 'wb') as fid: - fid.write(b64str_to_bytes(m_data)) + write_module_data(module_path, loaded_func_all['module_data']) return loaded_func_all['func'] -def get_function_data(job, internal_storage): +def _decode_data_byte_str(byte_str: Union[bytes, str]) -> bytes: + if isinstance(byte_str, bytes): + return byte_str + return ast.literal_eval(byte_str) + + +def get_function_data(job: SimpleNamespace, internal_storage) -> List[bytes]: """ - Get function data (iteradata) from storage + Gets the function data (iterdata) of every task of the job, either from + storage or from the invocation payload """ if job.data_key: extra_get_args = {} - if job.data_byte_ranges is not None: + if job.data_byte_ranges: init_byte = job.data_byte_ranges[0][0] last_byte = job.data_byte_ranges[-1][1] - range_str = f'bytes={init_byte}-{last_byte}' - extra_get_args['Range'] = range_str + extra_get_args['Range'] = f'bytes={init_byte}-{last_byte}' logger.info("Loading function data parameters from storage") - data_obj = internal_storage.get_data(job.data_key, extra_get_args=extra_get_args) + data_obj = internal_storage.get_data( + job.data_key, extra_get_args=extra_get_args + ) loaded_data = [] offset = 0 - if job.data_byte_ranges is not None: + if job.data_byte_ranges: for dbr in job.data_byte_ranges: length = dbr[1] - dbr[0] + 1 loaded_data.append(data_obj[offset:offset + length]) @@ -117,12 +114,14 @@ def get_function_data(job, internal_storage): else: loaded_data.append(data_obj) else: - loaded_data = [eval(byte_str) for byte_str in job.data_byte_strs] + loaded_data = [ + _decode_data_byte_str(byte_str) for byte_str in job.data_byte_strs + ] return loaded_data -def get_memory_usage(formatted=True): +def get_memory_usage(formatted: bool = True) -> Optional[Union[str, int]]: """ Gets the current memory usage of the runtime. To be used only in the action code. @@ -137,17 +136,17 @@ def get_memory_usage(formatted=True): discriminate_by_pid = False ps_mem.verify_environment(pids_to_show) - sorted_cmds, shareds, count, total, swaps, total_swap = \ - ps_mem.get_memory_usage(pids_to_show, split_args, discriminate_by_pid, - include_self=True, only_self=False) + _, _, _, total, _, _ = ps_mem.get_memory_usage( + pids_to_show, split_args, discriminate_by_pid, + include_self=True, only_self=False + ) if formatted: return sizeof_fmt(int(ps_mem.human(total, units=1))) - else: - return int(ps_mem.human(total, units=1)) + return int(ps_mem.human(total, units=1)) -def peak_memory(): - """Return the peak memory usage in bytes.""" +def peak_memory() -> Optional[int]: + """Returns the peak memory usage in bytes""" if not is_unix_system(): return None ru_maxrss = getrusage(RUSAGE_SELF).ru_maxrss @@ -156,7 +155,7 @@ def peak_memory(): return ru_maxrss * 1024 if platform.system() == "Linux" else ru_maxrss -def free_disk_space(dirname): +def free_disk_space(dirname: str) -> int: """ Returns the number of free bytes on the mount point containing DIRNAME """ @@ -164,65 +163,57 @@ def free_disk_space(dirname): return s.f_bsize * s.f_bavail -def get_server_info(): - """ - Returns server information - """ - container_name = subprocess.check_output("uname -n", shell=True).decode("ascii").strip() - ip_addr = subprocess.check_output("hostname -I", shell=True).decode("ascii").strip() - cores = subprocess.check_output("nproc", shell=True).decode("ascii").strip() - - cmd = "cat /sys/class/net/eth0/speed | awk '{print $0 / 1000\"GbE\"}'" - net_speed = subprocess.check_output(cmd, shell=True).decode("ascii").strip() +def _shell_output(cmd: str) -> str: + return subprocess.check_output(cmd, shell=True).decode("ascii").strip() - # cmd = "cat /sys/class/net/eth0/address" - # mac_address = subprocess.check_output(cmd, shell=True).decode("ascii").strip() - cmd = "grep MemTotal /proc/meminfo | awk '{print $2 / 1024 / 1024\"GB\"}'" - memory = subprocess.check_output(cmd, shell=True).decode("ascii").strip() - - server_info = {'container_name': container_name, - 'ip_address': ip_addr, - 'net_speed': net_speed, - 'cores': cores, - 'memory': memory} +def get_server_info() -> Dict[str, str]: """ - if os.path.exists("/proc"): - server_info.update({'/proc/cpuinfo': open("/proc/cpuinfo", 'r').read(), - '/proc/meminfo': open("/proc/meminfo", 'r').read(), - '/proc/self/cgroup': open("/proc/meminfo", 'r').read(), - '/proc/cgroups': open("/proc/cgroups", 'r').read()}) + Returns information about the machine this worker runs on """ - return server_info + net_speed_cmd = "cat /sys/class/net/eth0/speed | awk '{print $0 / 1000\"GbE\"}'" + memory_cmd = "grep MemTotal /proc/meminfo | awk '{print $2 / 1024 / 1024\"GB\"}'" + + return { + 'container_name': _shell_output("uname -n"), + 'ip_address': _shell_output("hostname -I"), + 'net_speed': _shell_output(net_speed_cmd), + 'cores': _shell_output("nproc"), + 'memory': _shell_output(memory_cmd), + } -def get_runtime_metadata(): +def get_runtime_metadata() -> Dict[str, Any]: """ Generates the runtime metadata needed for lithops """ - runtime_meta = dict() - mods = list(pkgutil.iter_modules()) - runtime_meta["preinstalls"] = [entry for entry in sorted([[mod, is_pkg] for _, mod, is_pkg in mods])] - python_version = sys.version_info - runtime_meta["python_version"] = str(python_version[0]) + "." + str(python_version[1]) - runtime_meta["lithops_version"] = lithops_ver + return { + "preinstalls": sorted( + [mod, is_pkg] for _, mod, is_pkg in pkgutil.iter_modules() + ), + "python_version": f"{sys.version_info[0]}.{sys.version_info[1]}", + "lithops_version": lithops_ver, + } - return runtime_meta - -def memory_monitor_worker(mm_conn, delay=0.01): +def memory_monitor_worker(mm_conn, delay: float = 0.01) -> None: """ - Monitor that checks the current memory usage + Monitors the memory usage of the runtime until the connection is ready, + and reports the peak back through it """ peak = 0 logger.debug("Starting memory monitor") + if get_memory_usage(formatted=False) is None: + # Nothing to measure here, so there is no point in polling + logger.debug("Memory monitor: memory usage is not available") + mm_conn.send(peak) + return + def make_measurement(peak): mem = get_memory_usage(formatted=False) + 5 * 1024**2 - if mem > peak: - peak = mem - return peak + return max(peak, mem) while not mm_conn.poll(delay): try: @@ -233,12 +224,13 @@ def make_measurement(peak): try: peak = make_measurement(peak) except Exception as e: - logger.error('Memory monitor: {}'.format(e)) + logger.error(f'Memory monitor: {e}') mm_conn.send(peak) @contextmanager def custom_redirection(fileobj): + """Redirects stdout and stderr to fileobj for the duration of the block""" old_stdout = sys.stdout old_stderr = sys.stderr sys.stdout = fileobj @@ -251,12 +243,17 @@ def custom_redirection(fileobj): class LogStream: + """ + Tees what the task prints to both the real stdout, so that it shows up in + the logs of the backend, and the task log file + """ def __init__(self, stream): self._stdout = sys.stdout self._stream = stream - def write(self, log): + def write(self, log: str) -> None: + """Writes to the log file, unless the handler closed it already""" self._stdout.write(log) try: self._stream.write(log) @@ -264,85 +261,98 @@ def write(self, log): except ValueError: pass - def flush(self): + def flush(self) -> None: + """Flushes both streams, unless the log file is closed already""" try: self._stream.flush() self._stdout.flush() except ValueError: pass - def fileno(self): + def fileno(self) -> int: + """Reports the descriptor of the real stdout""" return self._stdout.fileno() class SystemMonitor: + """ + Measures the resources that a process consumed between start() and + stop(). Monitors the current process if no process id is given + """ - def __init__(self, process_id=None): - """ - Initialize the SystemMonitor. - If process_id is None, monitor the current process. - """ + def __init__(self, process_id: Optional[int] = None): self.process_id = process_id self.cpu_usage = [] self.process = None self.cpu_times = None + self.start_net_io = None self.current_net_io = None self.mem_info = None - def start(self): + def start(self) -> None: """ - Start monitoring. + Starts monitoring, taking the baseline that stop() measures against """ if not psutil_found: return self.process = psutil.Process(self.process_id) - # record the initial CPU usage (to be ignored). + # The first measurement covers the whole life of the process, and is + # meant to be ignored: the next one is relative to this one psutil.cpu_percent(interval=None, percpu=True) - # Reset the network IO counters cache and baseline. + # psutil caches the counters, so they have to be cleared to get a + # fresh baseline psutil.net_io_counters.cache_clear() self.start_net_io = psutil.net_io_counters() - def stop(self): + def stop(self) -> None: """ - Stop monitoring. + Stops monitoring, recording everything consumed since start() """ if not psutil_found: return - # Record the CPU usage since the last call (start). self.cpu_usage = psutil.cpu_percent(interval=None, percpu=True) self.cpu_times = psutil.cpu_times() self.current_net_io = psutil.net_io_counters() self.mem_info = self.process.memory_full_info() - def get_cpu_info(self): + def get_cpu_info(self) -> Dict[str, Any]: """ - Return CPU usage, system time, and user time for each CPU core. + Returns the CPU usage of every core, and the system and user time """ if not psutil_found: return {"usage": [], "system": 0, "user": 0} - return {"usage": self.cpu_usage, "system": self.cpu_times.system, "user": self.cpu_times.user} + return { + "usage": self.cpu_usage, + "system": self.cpu_times.system, + "user": self.cpu_times.user, + } - def get_network_io(self): + def get_network_io(self) -> Dict[str, int]: """ - Calculate network IO (bytes sent and received) since the last reset. + Returns the bytes sent and received while monitoring """ if not psutil_found: return {"sent": 0, "recv": 0} - bytes_sent = self.current_net_io.bytes_sent - self.start_net_io.bytes_sent - bytes_recv = self.current_net_io.bytes_recv - self.start_net_io.bytes_recv - return {"sent": bytes_sent, "recv": bytes_recv} + return { + "sent": self.current_net_io.bytes_sent - self.start_net_io.bytes_sent, + "recv": self.current_net_io.bytes_recv - self.start_net_io.bytes_recv, + } - def get_memory_info(self): + def get_memory_info(self) -> Dict[str, int]: """ - Get memory usage information of the monitored process. + Returns the memory usage of the monitored process """ if not psutil_found: return {"rss": 0, "vms": 0, "uss": 0} - return {"rss": self.mem_info.rss, "vms": self.mem_info.vms, "uss": self.mem_info.uss} + return { + "rss": self.mem_info.rss, + "vms": self.mem_info.vms, + "uss": self.mem_info.uss, + } diff --git a/setup.py b/setup.py index 2c8b772ab..183efb73c 100644 --- a/setup.py +++ b/setup.py @@ -70,7 +70,8 @@ 'joblib': [ 'joblib', 'diskcache', - 'numpy' + 'numpy', + 'redis' ], 'plotting': [ 'pandas', From 0acf7967884c4a715b0c59f727c5786e277d0131 Mon Sep 17 00:00:00 2001 From: JosepSampe Date: Sat, 29 Aug 2026 14:50:28 +0200 Subject: [PATCH 2/4] Update CI deps --- lithops/tests/test_util.py | 68 +++++++++++++++++++++++--------------- setup.py | 12 +++++++ 2 files changed, 53 insertions(+), 27 deletions(-) diff --git a/lithops/tests/test_util.py b/lithops/tests/test_util.py index 476a6d01d..738613aab 100644 --- a/lithops/tests/test_util.py +++ b/lithops/tests/test_util.py @@ -20,12 +20,6 @@ import pytest -from lithops.util.ibm_token_manager import ( - COSTokenManager, - EXPIRY_MINUTES, - IAMTokenManager, - IBMTokenManager, -) from lithops.util.metrics import PrometheusExporter from lithops.util.ssh_client import SSHClient, ssh_boot_status_message @@ -179,28 +173,48 @@ def test_send_metric_swallows_post_errors(self, monkeypatch): exporter.send_metric('n', 1, type='gauge', labels=[]) -class _StubTokenManager(IBMTokenManager): - TOKEN_FILE = None - TYPE = 'TEST' - - def _generate_new_token(self): - self.token = 'new-token' - self.expiry_time = int( - (datetime.now(timezone.utc) + timedelta(hours=1)).timestamp() +class TestIBMTokenManager: + """ + ibm_token_manager imports ibm_botocore, which is only present with the + IBM extra. Skip this class when that extra is not installed so the rest + of the file still collects + """ + + @pytest.fixture(autouse=True) + def _ibm(self): + pytest.importorskip('ibm_botocore') + pytest.importorskip('ibm_cloud_sdk_core') + from lithops.util.ibm_token_manager import ( + COSTokenManager, + EXPIRY_MINUTES, + IAMTokenManager, + IBMTokenManager, ) + class StubTokenManager(IBMTokenManager): + TOKEN_FILE = None + TYPE = 'TEST' -class TestIBMTokenManager: + def _generate_new_token(self): + self.token = 'new-token' + self.expiry_time = int( + (datetime.now(timezone.utc) + timedelta(hours=1)).timestamp() + ) + + self.COSTokenManager = COSTokenManager + self.EXPIRY_MINUTES = EXPIRY_MINUTES + self.IAMTokenManager = IAMTokenManager + self.StubTokenManager = StubTokenManager def test_token_file_constant_is_spelled_correctly(self): - assert hasattr(COSTokenManager, 'TOKEN_FILE') - assert hasattr(IAMTokenManager, 'TOKEN_FILE') - assert 'ibm_cos' in COSTokenManager.TOKEN_FILE - assert 'ibm_iam' in IAMTokenManager.TOKEN_FILE - assert not hasattr(COSTokenManager, 'TOEKN_FILE') + assert hasattr(self.COSTokenManager, 'TOKEN_FILE') + assert hasattr(self.IAMTokenManager, 'TOKEN_FILE') + assert 'ibm_cos' in self.COSTokenManager.TOKEN_FILE + assert 'ibm_iam' in self.IAMTokenManager.TOKEN_FILE + assert not hasattr(self.COSTokenManager, 'TOEKN_FILE') def test_missing_expiry_is_expired(self): - mgr = _StubTokenManager('key') + mgr = self.StubTokenManager('key') assert mgr._get_token_minutes_left() == 0 assert mgr._is_token_expired() @@ -208,16 +222,16 @@ def test_reuses_unexpired_token(self): expiry = int( (datetime.now(timezone.utc) + timedelta(hours=2)).timestamp() ) - mgr = _StubTokenManager('key', token='cached', token_expiry_time=expiry) - assert mgr._get_token_minutes_left() >= EXPIRY_MINUTES + mgr = self.StubTokenManager('key', token='cached', token_expiry_time=expiry) + assert mgr._get_token_minutes_left() >= self.EXPIRY_MINUTES token, exp = mgr.get_token() assert token == 'cached' assert exp == expiry def test_refresh_dumps_and_returns_new_token(self, tmp_path, monkeypatch): path = tmp_path / 'token' - monkeypatch.setattr(_StubTokenManager, 'TOKEN_FILE', str(path)) - mgr = _StubTokenManager('key') + monkeypatch.setattr(self.StubTokenManager, 'TOKEN_FILE', str(path)) + mgr = self.StubTokenManager('key') with patch( 'lithops.util.ibm_token_manager.dump_yaml_config' ) as dump: @@ -229,7 +243,7 @@ def test_refresh_dumps_and_returns_new_token(self, tmp_path, monkeypatch): def test_loads_cache_file(self, tmp_path, monkeypatch): path = tmp_path / 'token' - monkeypatch.setattr(_StubTokenManager, 'TOKEN_FILE', str(path)) + monkeypatch.setattr(self.StubTokenManager, 'TOKEN_FILE', str(path)) path.write_text('x') expiry = int( (datetime.now(timezone.utc) + timedelta(hours=2)).timestamp() @@ -242,7 +256,7 @@ def test_loads_cache_file(self, tmp_path, monkeypatch): 'lithops.util.ibm_token_manager.os.path.exists', return_value=True, ): - mgr = _StubTokenManager('key') + mgr = self.StubTokenManager('key') assert mgr.token == 'from-disk' assert mgr.expiry_time == expiry diff --git a/setup.py b/setup.py index 183efb73c..7a180fb46 100644 --- a/setup.py +++ b/setup.py @@ -86,6 +86,18 @@ 'pytest', 'kubernetes', 'pika', + 'ibm-cos-sdk', + 'ibm-cloud-sdk-core', + 'joblib', + 'diskcache', + 'numpy', + 'redis', + 'pandas', + 'matplotlib', + 'seaborn', + 'flask', + 'gevent', + 'scikit-learn', ] } From b0caca7d6cef66f4247534405ef5c775105f6bd3 Mon Sep 17 00:00:00 2001 From: JosepSampe Date: Sat, 29 Aug 2026 15:11:14 +0200 Subject: [PATCH 3/4] Update tests --- lithops/tests/test_monitor.py | 1 + lithops/tests/test_wait.py | 32 ++++++++++++++++++-------------- 2 files changed, 19 insertions(+), 14 deletions(-) diff --git a/lithops/tests/test_monitor.py b/lithops/tests/test_monitor.py index 18d27e0b8..f4776fe0e 100644 --- a/lithops/tests/test_monitor.py +++ b/lithops/tests/test_monitor.py @@ -323,6 +323,7 @@ def test_print_status_log_does_not_repeat_when_all_finished(self, caplog): monitor = _monitor() monitor.add_futures([FakeFuture('M000', invoked=True, ready=True)]) first, _ = monitor._print_status_log(previous_log=None, log_time=0) + caplog.clear() with caplog.at_level(logging.DEBUG, logger='lithops.monitor'): counts, log_time = monitor._print_status_log( previous_log=first, log_time=LOG_INTERVAL + 1 diff --git a/lithops/tests/test_wait.py b/lithops/tests/test_wait.py index 561be768b..c19a9aee3 100644 --- a/lithops/tests/test_wait.py +++ b/lithops/tests/test_wait.py @@ -12,6 +12,7 @@ # limitations under the License. # +import importlib import signal import threading from types import SimpleNamespace @@ -22,6 +23,8 @@ from lithops.utils import is_unix_system from lithops.utils import FuturesList + +wait_mod = importlib.import_module('lithops.wait') from lithops.wait import ( ALL_COMPLETED, ALWAYS, @@ -144,7 +147,7 @@ def test_wraps_single_complete_future(self): class TestCreateExecutorsData: - @patch('lithops.wait.InternalStorage') + @patch.object(wait_mod, 'InternalStorage') def test_groups_futures_and_reuses_matching_storage(self, mock_storage_cls): internal = MagicMock() internal.backend = 'localhost' @@ -238,7 +241,7 @@ def test_always_polls_once_without_looping(self): monitor.storage_backend = 'localhost' internal = MagicMock() internal.backend = 'localhost' - with patch('lithops.wait._get_executor_data', return_value=0) as get: + with patch.object(wait_mod, '_get_executor_data', return_value=0) as get: wait( [future], return_when=ALWAYS, @@ -250,8 +253,9 @@ def test_always_polls_once_without_looping(self): def test_keyboard_interrupt_reraises_after_logging(self): future = FakeFuture() - with patch( - 'lithops.wait._create_executors_data_from_futures', + with patch.object( + wait_mod, + '_create_executors_data_from_futures', side_effect=KeyboardInterrupt, ): with pytest.raises(KeyboardInterrupt): @@ -270,9 +274,9 @@ def get_data(fs, exec_data, **kwargs): future.success = True return 1 - with patch('lithops.wait.JobMonitor', return_value=monitor) as cls, \ - patch('lithops.wait._get_executor_data', side_effect=get_data), \ - patch('lithops.wait.time.sleep'): + with patch.object(wait_mod, 'JobMonitor', return_value=monitor) as cls, \ + patch.object(wait_mod, '_get_executor_data', side_effect=get_data), \ + patch.object(wait_mod.time, 'sleep'): wait( [future], show_progressbar=False, @@ -301,10 +305,10 @@ def get_data(fs, exec_data, **kwargs): future.success = True return 1 - with patch('lithops.wait.signal.signal', side_effect=fake_signal), \ - patch('lithops.wait.signal.alarm') as alarm, \ - patch('lithops.wait._get_executor_data', side_effect=get_data), \ - patch('lithops.wait.time.sleep'): + with patch.object(wait_mod.signal, 'signal', side_effect=fake_signal), \ + patch.object(wait_mod.signal, 'alarm') as alarm, \ + patch.object(wait_mod, '_get_executor_data', side_effect=get_data), \ + patch.object(wait_mod.time, 'sleep'): wait( [future], timeout=17, @@ -342,8 +346,8 @@ def sleep(seconds): if threading.current_thread() is test_thread: sleeps.append(seconds) - with patch('lithops.wait._get_executor_data', side_effect=get_data), \ - patch('lithops.wait.time.sleep', side_effect=sleep): + with patch.object(wait_mod, '_get_executor_data', side_effect=get_data), \ + patch.object(wait_mod.time, 'sleep', side_effect=sleep): wait( [future], return_when=ALL_COMPLETED, @@ -373,7 +377,7 @@ def parent_status(**kwargs): internal = MagicMock() internal.backend = 'localhost' - with patch('lithops.wait.time.sleep'): + with patch.object(wait_mod.time, 'sleep'): done, not_done = wait( [parent], show_progressbar=False, From aa319066a55072f82e8a5731278e0acac2ff4ce4 Mon Sep 17 00:00:00 2001 From: JosepSampe Date: Sat, 29 Aug 2026 15:17:08 +0200 Subject: [PATCH 4/4] Fix linting --- lithops/tests/test_joblib.py | 36 +++++++++++++++++++++++------------- lithops/tests/test_wait.py | 8 +++----- 2 files changed, 26 insertions(+), 18 deletions(-) diff --git a/lithops/tests/test_joblib.py b/lithops/tests/test_joblib.py index a3ed4a169..145972947 100644 --- a/lithops/tests/test_joblib.py +++ b/lithops/tests/test_joblib.py @@ -128,6 +128,20 @@ def _sum_with(shared, i): return int(shared.sum()) + i +def _tiny_classification_data(n=40, n_features=8, n_classes=3): + """ + A small labelled set for the sklearn searches. load_digits() assigns + to ndarray.shape, which NumPy 2.5 warns on + """ + import numpy as np + + rng = np.random.default_rng(0) + return ( + rng.normal(size=(n, n_features)), + rng.integers(0, n_classes, size=n), + ) + + class TestSklearnOverJoblib: """ The searches the examples in examples/ run, in a smaller shape so that @@ -139,11 +153,10 @@ def _needs_sklearn(self): pytest.importorskip('sklearn') def test_grid_search_over_the_lithops_backend(self): - from sklearn.datasets import load_digits from sklearn.model_selection import GridSearchCV from sklearn.tree import DecisionTreeClassifier - digits = load_digits() + X, y = _tiny_classification_data() search = GridSearchCV( DecisionTreeClassifier(random_state=0), {'max_depth': [2, 4]}, @@ -152,12 +165,12 @@ def test_grid_search_over_the_lithops_backend(self): ) with _on_localhost(): - search.fit(digits.data, digits.target) + search.fit(X, y) assert search.best_params_['max_depth'] in (2, 4) assert 0.0 < search.best_score_ <= 1.0 # refit ran, so the search can predict - assert len(search.predict(digits.data[:5])) == 5 + assert len(search.predict(X[:5])) == 5 def test_the_dataset_is_proxied_for_every_fit_of_the_search(self): # Every fit gets the same X and y, so they travel as one cloud object @@ -165,13 +178,12 @@ def test_the_dataset_is_proxied_for_every_fit_of_the_search(self): # carries a fourth element with their positions from unittest.mock import patch - from sklearn.datasets import load_digits from sklearn.model_selection import GridSearchCV from sklearn.tree import DecisionTreeClassifier from lithops.util.joblib import lithops_backend - digits = load_digits() + X, y = _tiny_classification_data() search = GridSearchCV( DecisionTreeClassifier(random_state=0), {'max_depth': [2, 4, 6]}, @@ -184,7 +196,7 @@ def test_the_dataset_is_proxied_for_every_fit_of_the_search(self): _counting_optimizer(proxied) ): with _on_localhost(): - search.fit(digits.data, digits.target) + search.fit(X, y) assert proxied, 'the batch never went through the optimizer' # Three candidates over two folds @@ -192,11 +204,10 @@ def test_the_dataset_is_proxied_for_every_fit_of_the_search(self): def test_randomized_search_over_the_lithops_backend(self): import numpy as np - from sklearn.datasets import load_digits from sklearn.model_selection import RandomizedSearchCV from sklearn.tree import DecisionTreeClassifier - digits = load_digits() + X, y = _tiny_classification_data() search = RandomizedSearchCV( DecisionTreeClassifier(random_state=0), {'min_samples_leaf': np.arange(1, 10)}, @@ -206,19 +217,18 @@ def test_randomized_search_over_the_lithops_backend(self): ) with _on_localhost(): - search.fit(digits.data, digits.target) + search.fit(X, y) assert 0.0 < search.best_score_ <= 1.0 def test_a_pipeline_search_over_the_lithops_backend(self): # The shape of examples/sklearn_job_3.py, without pandas - from sklearn.datasets import load_digits from sklearn.model_selection import GridSearchCV from sklearn.pipeline import Pipeline from sklearn.preprocessing import StandardScaler from sklearn.tree import DecisionTreeClassifier - digits = load_digits() + X, y = _tiny_classification_data() pipeline = Pipeline([ ('scale', StandardScaler()), ('classifier', DecisionTreeClassifier(random_state=0)), @@ -228,7 +238,7 @@ def test_a_pipeline_search_over_the_lithops_backend(self): ) with _on_localhost(): - search.fit(digits.data, digits.target) + search.fit(X, y) assert search.best_params_['classifier__max_depth'] in (2, 4) assert 0.0 < search.best_score_ <= 1.0 diff --git a/lithops/tests/test_wait.py b/lithops/tests/test_wait.py index c19a9aee3..2de70d8ec 100644 --- a/lithops/tests/test_wait.py +++ b/lithops/tests/test_wait.py @@ -20,11 +20,7 @@ import pytest -from lithops.utils import is_unix_system - -from lithops.utils import FuturesList - -wait_mod = importlib.import_module('lithops.wait') +from lithops.utils import FuturesList, is_unix_system from lithops.wait import ( ALL_COMPLETED, ALWAYS, @@ -42,6 +38,8 @@ wait, ) +wait_mod = importlib.import_module('lithops.wait') + class FakeFuture: def __init__(self, *, done=False, success=False, ready=False, executor_id='sess-0',