Source code for cyto.runners.container_worker

import sys
import cloudpickle



def main():
    """
    A generic worker for executing tasks inside a container.

    This script is not meant to be called directly by the user. It is called by
    the runner system (e.g., DockerRunner) inside a container. It deserializes a
    task object and its input data, executes the task's baremetal method, and
    serializes the result to a file.

    Args:
        task_path (str): Path to the pickled task object.
        data_path (str): Path to the pickled input data.
        result_path (str): Path where the pickled result should be saved.
    """
    if len(sys.argv) != 4:
        print("Usage: python -m cyto.runners.container_worker <task_path> <data_path> <result_path>", file=sys.stderr)
        sys.exit(1)

    _, task_path, data_path, result_path = sys.argv

    # Deserialize task and data
    try:
        with open(task_path, "rb") as f:
            task = cloudpickle.load(f)
        with open(data_path, "rb") as f:
            data = cloudpickle.load(f)
    except FileNotFoundError as e:
        print(f"Error: Input file not found - {e}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"Error deserializing task or data: {e}", file=sys.stderr)
        sys.exit(1)


    # Execute the task
    # Inside the container, we always run the baremetal version of the task.
    try:
        result = task.run_baremetal(data)
    except Exception as e:
        print(f"Error during task execution: {e}", file=sys.stderr)
        sys.exit(1)


    # Serialize the result
    try:
        with open(result_path, "wb") as f:
            cloudpickle.dump(result, f)
    except Exception as e:
        print(f"Error serializing result: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()