runnable 0.34.0a1__py3-none-any.whl → 1.0.0__py3-none-any.whl

extensions/nodes/torch.py

@@ -1,273 +0,0 @@
-import importlib
-import logging
-import os
-import random
-import string
-from datetime import datetime
-from pathlib import Path
-from typing import Any, Callable, Optional
-
-from pydantic import BaseModel, ConfigDict, Field, field_serializer
-
-from extensions.nodes.torch_config import EasyTorchConfig, TorchConfig
-from runnable import PythonJob, datastore, defaults
-from runnable.datastore import StepLog
-from runnable.nodes import ExecutableNode
-from runnable.tasks import PythonTaskType, create_task
-from runnable.utils import TypeMapVariable
-
-logger = logging.getLogger(defaults.LOGGER_NAME)
-
-try:
-    from torch.distributed.elastic.multiprocessing.api import DefaultLogsSpecs, Std
-    from torch.distributed.launcher.api import LaunchConfig, elastic_launch
-except ImportError:
-    logger.exception("Torch is not installed. Please install torch first.")
-    raise Exception("Torch is not installed. Please install torch first.")
-
-
-def training_subprocess():
-    """
-    This function is called by the torch.distributed.launcher.api.elastic_launch
-    It happens in a subprocess and is responsible for executing the user's function
-
-    It is unrelated to the actual node execution, so any cataloging, run_log_store should be
-    handled to match to main process.
-
-    We have these variables to use:
-
-    os.environ["RUNNABLE_TORCH_COMMAND"] = self.executable.command
-    os.environ["RUNNABLE_TORCH_PARAMETERS_FILES"] = (
-        self._context.parameters_file or ""
-    )
-    os.environ["RUNNABLE_TORCH_RUN_ID"] = self._context.run_id
-    os.environ["RUNNABLE_TORCH_COPY_CONTENTS_TO"] = (
-        self._context.catalog_handler.compute_data_folder
-    )
-    os.environ["RUNNABLE_TORCH_TORCH_LOGS"] = self.log_dir or ""
-
-    """
-    command = os.environ.get("RUNNABLE_TORCH_COMMAND")
-    run_id = os.environ.get("RUNNABLE_TORCH_RUN_ID", "")
-    parameters_files = os.environ.get("RUNNABLE_TORCH_PARAMETERS_FILES", "")
-
-    process_run_id = (
-        run_id
-        + "-"
-        + os.environ.get("RANK", "")
-        + "-"
-        + "".join(random.choices(string.ascii_lowercase, k=3))
-    )
-    os.environ["TORCH_DISTRIBUTED_DEBUG"] = "DETAIL"
-
-    delete_env_vars_with_prefix("RUNNABLE_")
-
-    func = get_callable_from_dotted_path(command)
-
-    # The job runs with the default configuration
-    # ALl the execution logs are stored in .catalog
-    job = PythonJob(function=func)
-
-    job.execute(
-        parameters_file=parameters_files,
-        job_id=process_run_id,
-    )
-
-    from runnable.context import run_context
-
-    job_log = run_context.run_log_store.get_run_log_by_id(run_id=run_context.run_id)
-
-    if job_log.status == defaults.FAIL:
-        raise Exception(f"Job {process_run_id} failed")
-
-
-# TODO: Can this be utils.get_module_and_attr_names
-def get_callable_from_dotted_path(dotted_path) -> Callable:
-    try:
-        # Split the path into module path and callable object
-        module_path, callable_name = dotted_path.rsplit(".", 1)
-
-        # Import the module
-        module = importlib.import_module(module_path)
-
-        # Get the callable from the module
-        callable_obj = getattr(module, callable_name)
-
-        # Check if the object is callable
-        if not callable(callable_obj):
-            raise TypeError(f"The object {callable_name} is not callable.")
-
-        return callable_obj
-
-    except (ImportError, AttributeError, ValueError) as e:
-        raise ImportError(f"Could not import '{dotted_path}'.") from e
-
-
-def delete_env_vars_with_prefix(prefix):
-    to_delete = []  # List to keep track of variables to delete
-
-    # Iterate over a list of all environment variable keys
-    for var in os.environ:
-        if var.startswith(prefix):
-            to_delete.append(var)
-
-    # Delete each of the variables collected
-    for var in to_delete:
-        del os.environ[var]
-
-
-# TODO: The design of this class is not final
-class TorchNode(ExecutableNode, TorchConfig):
-    node_type: str = Field(default="torch", serialization_alias="type")
-    executable: PythonTaskType = Field(exclude=True)
-
-    # Similar to TaskNode
-    model_config = ConfigDict(extra="allow")
-
-    def get_summary(self) -> dict[str, Any]:
-        summary = {
-            "name": self.name,
-            "type": self.node_type,
-        }
-
-        return summary
-
-    @classmethod
-    def parse_from_config(cls, config: dict[str, Any]) -> "TorchNode":
-        task_config = {
-            k: v for k, v in config.items() if k not in TorchNode.model_fields.keys()
-        }
-        node_config = {
-            k: v for k, v in config.items() if k in TorchNode.model_fields.keys()
-        }
-
-        executable = create_task(task_config)
-
-        assert isinstance(executable, PythonTaskType)
-        return cls(executable=executable, **node_config, **task_config)
-
-    def get_launch_config(self) -> LaunchConfig:
-        internal_log_spec = InternalLogSpecs(**self.model_dump(exclude_none=True))
-        log_spec: DefaultLogsSpecs = DefaultLogsSpecs(
-            **internal_log_spec.model_dump(exclude_none=True)
-        )
-        easy_torch_config = EasyTorchConfig(
-            **self.model_dump(
-                exclude_none=True,
-            )
-        )
-
-        launch_config = LaunchConfig(
-            **easy_torch_config.model_dump(
-                exclude_none=True,
-            ),
-            logs_specs=log_spec,
-            run_id=self._context.run_id,
-        )
-        logger.info(f"launch_config: {launch_config}")
-        return launch_config
-
-    def execute(
-        self,
-        mock=False,
-        map_variable: TypeMapVariable = None,
-        attempt_number: int = 1,
-    ) -> StepLog:
-        assert (
-            map_variable is None or not map_variable
-        ), "TorchNode does not support map_variable"
-
-        step_log = self._context.run_log_store.get_step_log(
-            self._get_step_log_name(map_variable), self._context.run_id
-        )
-
-        # Attempt to call the function or elastic launch
-        launch_config = self.get_launch_config()
-        logger.info(f"launch_config: {launch_config}")
-
-        # ENV variables are shared with the subprocess, use that as communication
-        os.environ["RUNNABLE_TORCH_COMMAND"] = self.executable.command
-        os.environ["RUNNABLE_TORCH_PARAMETERS_FILES"] = (
-            self._context.parameters_file or ""
-        )
-        os.environ["RUNNABLE_TORCH_RUN_ID"] = self._context.run_id
-
-        launcher = elastic_launch(
-            launch_config,
-            training_subprocess,
-        )
-        try:
-            launcher()
-            attempt_log = datastore.StepAttempt(
-                status=defaults.SUCCESS,
-                start_time=str(datetime.now()),
-                end_time=str(datetime.now()),
-                attempt_number=attempt_number,
-            )
-        except Exception as e:
-            attempt_log = datastore.StepAttempt(
-                status=defaults.FAIL,
-                start_time=str(datetime.now()),
-                end_time=str(datetime.now()),
-                attempt_number=attempt_number,
-            )
-            logger.error(f"Error executing TorchNode: {e}")
-        finally:
-            # This can only come from the subprocess
-            if Path(".catalog").exists():
-                os.rename(".catalog", "proc_logs")
-                # Move .catalog and torch_logs to the parent node's catalog location
-                self._context.catalog_handler.put(
-                    "proc_logs/**/*", allow_file_not_found_exc=True
-                )
-
-            # TODO: This is not working!!
-            if self.log_dir:
-                self._context.catalog_handler.put(
-                    self.log_dir + "/**/*", allow_file_not_found_exc=True
-                )
-
-        delete_env_vars_with_prefix("RUNNABLE_TORCH")
-
-        logger.info(f"attempt_log: {attempt_log}")
-        logger.info(f"Step {self.name} completed with status: {attempt_log.status}")
-
-        step_log.status = attempt_log.status
-        step_log.attempts.append(attempt_log)
-
-        return step_log
-
-    def fan_in(self, map_variable: dict[str, str | int | float] | None = None):
-        # Destroy the service
-        # Destroy the statefulset
-        assert (
-            map_variable is None or not map_variable
-        ), "TorchNode does not support map_variable"
-
-    def fan_out(self, map_variable: dict[str, str | int | float] | None = None):
-        # Create a service
-        # Create a statefulset
-        # Gather the IPs and set them as parameters downstream
-        assert (
-            map_variable is None or not map_variable
-        ), "TorchNode does not support map_variable"
-
-
-# This internal model makes it easier to extract the required fields
-# of log specs from user specification.
-# https://github.com/pytorch/pytorch/blob/main/torch/distributed/elastic/multiprocessing/api.py#L243
-class InternalLogSpecs(BaseModel):
-    log_dir: Optional[str] = Field(default="torch_logs")
-    redirects: str = Field(default="0")  # Std.NONE
-    tee: str = Field(default="0")  # Std.NONE
-    local_ranks_filter: Optional[set[int]] = Field(default=None)
-
-    model_config = ConfigDict(extra="ignore")
-
-    @field_serializer("redirects")
-    def convert_redirects(self, redirects: str) -> Std | dict[int, Std]:
-        return Std.from_str(redirects)
-
-    @field_serializer("tee")
-    def convert_tee(self, tee: str) -> Std | dict[int, Std]:
-        return Std.from_str(tee)

extensions/nodes/torch_config.py

@@ -1,76 +0,0 @@
-from enum import Enum
-from typing import Any, Optional
-
-from pydantic import BaseModel, ConfigDict, Field, computed_field
-
-
-class StartMethod(str, Enum):
-    spawn = "spawn"
-    fork = "fork"
-    forkserver = "forkserver"
-
-
-## The idea is the following:
-# Users can configure any of the options present in TorchConfig class.
-# The LaunchConfig class will be created from TorchConfig.
-# The LogSpecs is sent as a parameter to the launch config.
-
-## NO idea of standalone and how to send it
-
-
-# The user sees this as part of the config of the node.
-# It is kept as similar as possible to torchrun
-class TorchConfig(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    # excluded as LaunchConfig requires min and max nodes
-    nnodes: str = Field(default="1:1", exclude=True, description="min:max")
-    nproc_per_node: int = Field(default=1, description="Number of processes per node")
-
-    # will be used to create the log specs
-    # But they are excluded from dump as logs specs is a class for LaunchConfig
-    # from_str("0") -> Std.NONE
-    # from_str("1") -> Std.OUT
-    # from_str("0:3,1:0,2:1,3:2") -> {0: Std.ALL, 1: Std.NONE, 2: Std.OUT, 3: Std.ERR}
-    log_dir: Optional[str] = Field(default="torch_logs", exclude=True)
-    redirects: str = Field(default="0", exclude=True)  # Std.NONE
-    tee: str = Field(default="0", exclude=True)  # Std.NONE
-    local_ranks_filter: Optional[set[int]] = Field(default=None, exclude=True)
-
-    role: str | None = Field(default=None)
-
-    # run_id would be the run_id of the context
-    # and sent at the creation of the LaunchConfig
-
-    # This section is about the communication between nodes/processes
-    rdzv_backend: str | None = Field(default="static")
-    rdzv_endpoint: str | None = Field(default="")
-    rdzv_configs: dict[str, Any] = Field(default_factory=dict)
-    rdzv_timeout: int | None = Field(default=None)
-
-    max_restarts: int | None = Field(default=None)
-    monitor_interval: float | None = Field(default=None)
-    start_method: str | None = Field(default=StartMethod.spawn)
-    log_line_prefix_template: str | None = Field(default=None)
-    local_addr: Optional[str] = None
-
-    # https://github.com/pytorch/pytorch/blob/main/torch/distributed/run.py#L753
-    # master_addr: str | None = Field(default="localhost")
-    # master_port: str | None = Field(default="29500")
-    # training_script: str = Field(default="dummy_training_script")
-    # training_script_args: str = Field(default="")
-
-
-class EasyTorchConfig(TorchConfig):
-    model_config = ConfigDict(extra="ignore")
-
-    # TODO: Validate min < max
-    @computed_field  # type: ignore
-    @property
-    def min_nodes(self) -> int:
-        return int(self.nnodes.split(":")[0])
-
-    @computed_field  # type: ignore
-    @property
-    def max_nodes(self) -> int:
-        return int(self.nnodes.split(":")[1])

extensions/tasks/torch.py

@@ -1,286 +0,0 @@
-import importlib
-import logging
-import os
-import random
-import string
-from datetime import datetime
-from pathlib import Path
-from typing import Any, Optional
-
-from pydantic import BaseModel, ConfigDict, Field, field_serializer, model_validator
-from ruamel.yaml import YAML
-
-import runnable.context as context
-from extensions.tasks.torch_config import EasyTorchConfig, TorchConfig
-from runnable import Catalog, defaults
-from runnable.datastore import StepAttempt
-from runnable.tasks import BaseTaskType
-from runnable.utils import get_module_and_attr_names
-
-logger = logging.getLogger(defaults.LOGGER_NAME)
-
-logger = logging.getLogger(defaults.LOGGER_NAME)
-
-try:
-    from torch.distributed.elastic.multiprocessing.api import DefaultLogsSpecs, Std
-    from torch.distributed.launcher.api import LaunchConfig, elastic_launch
-
-except ImportError as e:
-    logger.exception("torch is not installed")
-    raise Exception("torch is not installed") from e
-
-
-def get_min_max_nodes(nnodes: str) -> tuple[int, int]:
-    min_nodes, max_nodes = (int(x) for x in nnodes.split(":"))
-    return min_nodes, max_nodes
-
-
-class TorchTaskType(BaseTaskType, TorchConfig):
-    task_type: str = Field(default="torch", serialization_alias="command_type")
-    catalog: Optional[Catalog] = Field(default=None, alias="catalog")
-    command: str
-
-    @model_validator(mode="before")
-    @classmethod
-    def check_secrets_and_returns(cls, data: Any) -> Any:
-        if isinstance(data, dict):
-            if "secrets" in data and data["secrets"]:
-                raise ValueError("'secrets' is not supported for torch")
-            if "returns" in data and data["returns"]:
-                raise ValueError("'secrets' is not supported for torch")
-        return data
-
-    def get_summary(self) -> dict[str, Any]:
-        return self.model_dump(by_alias=True, exclude_none=True)
-
-    @property
-    def _context(self):
-        return context.run_context
-
-    def _get_launch_config(self) -> LaunchConfig:
-        internal_log_spec = InternalLogSpecs(**self.model_dump(exclude_none=True))
-        log_spec: DefaultLogsSpecs = DefaultLogsSpecs(
-            **internal_log_spec.model_dump(exclude_none=True)
-        )
-        easy_torch_config = EasyTorchConfig(
-            **self.model_dump(
-                exclude_none=True,
-            )
-        )
-        print("###", easy_torch_config)
-        print("###", easy_torch_config)
-        launch_config = LaunchConfig(
-            **easy_torch_config.model_dump(
-                exclude_none=True,
-            ),
-            logs_specs=log_spec,
-            run_id=self._context.run_id,
-        )
-        logger.info(f"launch_config: {launch_config}")
-        return launch_config
-
-    def execute_command(
-        self,
-        map_variable: defaults.TypeMapVariable = None,
-    ):
-        assert map_variable is None, "map_variable is not supported for torch"
-
-        # The below should happen only if we are in the node that we want to execute
-        # For a single node, multi worker setup, this should be the entry point
-        # For a multi-node, we need to:
-        # - create a service config
-        # - Create a stateful set with number of nodes
-        # - Create a job to run the torch.distributed.launcher.api.elastic_launch on every node
-        # - the entry point to runnnable could be a way to trigger execution instead of scaling
-        is_execute = os.environ.get("RUNNABLE_TORCH_EXECUTE", "true") == "true"
-
-        _, max_nodes = get_min_max_nodes(self.nnodes)
-
-        if max_nodes > 1 and not is_execute:
-            executor = self._context.executor
-            executor.scale_up(self)
-            return StepAttempt(
-                status=defaults.SUCCESS,
-                start_time=str(datetime.now()),
-                end_time=str(datetime.now()),
-                attempt_number=1,
-                message="Triggered a scale up",
-            )
-
-        # The below should happen only if we are in the node that we want to execute
-        # For a single node, multi worker setup, this should be the entry point
-        # For a multi-node, we need to:
-        # - create a service config
-        # - Create a stateful set with number of nodes
-        # - Create a job to run the torch.distributed.launcher.api.elastic_launch on every node
-        # - the entry point to runnnable could be a way to trigger execution instead of scaling
-        is_execute = os.environ.get("RUNNABLE_TORCH_EXECUTE", "true") == "true"
-
-        _, max_nodes = get_min_max_nodes(self.nnodes)
-
-        if max_nodes > 1 and not is_execute:
-            executor = self._context.executor
-            executor.scale_up(self)
-            return StepAttempt(
-                status=defaults.SUCCESS,
-                start_time=str(datetime.now()),
-                end_time=str(datetime.now()),
-                attempt_number=1,
-                message="Triggered a scale up",
-            )
-
-        launch_config = self._get_launch_config()
-        print("###****", launch_config)
-        print("###****", launch_config)
-        logger.info(f"launch_config: {launch_config}")
-
-        # ENV variables are shared with the subprocess, use that as communication
-        os.environ["RUNNABLE_TORCH_COMMAND"] = self.command
-        os.environ["RUNNABLE_TORCH_PARAMETERS_FILES"] = (
-            self._context.parameters_file or ""
-        )
-        os.environ["RUNNABLE_TORCH_RUN_ID"] = self._context.run_id
-
-        launcher = elastic_launch(
-            launch_config,
-            training_subprocess,
-        )
-        try:
-            launcher()
-            attempt_log = StepAttempt(
-                status=defaults.SUCCESS,
-                start_time=str(datetime.now()),
-                end_time=str(datetime.now()),
-                attempt_number=1,
-            )
-        except Exception as e:
-            attempt_log = StepAttempt(
-                status=defaults.FAIL,
-                start_time=str(datetime.now()),
-                end_time=str(datetime.now()),
-                attempt_number=1,
-            )
-            logger.error(f"Error executing TorchNode: {e}")
-        finally:
-            # This can only come from the subprocess
-            if Path("proc_logs").exists():
-                # Move .catalog and torch_logs to the parent node's catalog location
-                self._context.catalog_handler.put(
-                    "proc_logs/**/*", allow_file_not_found_exc=True
-                )
-
-            # TODO: This is not working!!
-            if self.log_dir:
-                self._context.catalog_handler.put(
-                    self.log_dir + "/**/*", allow_file_not_found_exc=True
-                )
-
-        delete_env_vars_with_prefix("RUNNABLE_TORCH")
-        logger.info(f"attempt_log: {attempt_log}")
-
-        return attempt_log
-
-
-# This internal model makes it easier to extract the required fields
-# of log specs from user specification.
-# https://github.com/pytorch/pytorch/blob/main/torch/distributed/elastic/multiprocessing/api.py#L243
-class InternalLogSpecs(BaseModel):
-    log_dir: Optional[str] = Field(default="torch_logs")
-    redirects: str = Field(default="0")  # Std.NONE
-    tee: str = Field(default="0")  # Std.NONE
-    local_ranks_filter: Optional[set[int]] = Field(default=None)
-
-    model_config = ConfigDict(extra="ignore")
-
-    @field_serializer("redirects")
-    def convert_redirects(self, redirects: str) -> Std | dict[int, Std]:
-        return Std.from_str(redirects)
-
-    @field_serializer("tee")
-    def convert_tee(self, tee: str) -> Std | dict[int, Std]:
-        return Std.from_str(tee)
-
-
-def delete_env_vars_with_prefix(prefix):
-    to_delete = []  # List to keep track of variables to delete
-
-    # Iterate over a list of all environment variable keys
-    for var in os.environ:
-        if var.startswith(prefix):
-            to_delete.append(var)
-
-    # Delete each of the variables collected
-    for var in to_delete:
-        del os.environ[var]
-
-
-def training_subprocess():
-    """
-    This function is called by the torch.distributed.launcher.api.elastic_launch
-    It happens in a subprocess and is responsible for executing the user's function
-
-    It is unrelated to the actual node execution, so any cataloging, run_log_store should be
-    handled to match to main process.
-
-    We have these variables to use:
-
-    os.environ["RUNNABLE_TORCH_COMMAND"] = self.executable.command
-    os.environ["RUNNABLE_TORCH_PARAMETERS_FILES"] = (
-        self._context.parameters_file or ""
-    )
-    os.environ["RUNNABLE_TORCH_RUN_ID"] = self._context.run_id
-    os.environ["RUNNABLE_TORCH_TORCH_LOGS"] = self.log_dir or ""
-
-    """
-    from runnable import PythonJob  # noqa: F401
-
-    command = os.environ.get("RUNNABLE_TORCH_COMMAND")
-    assert command, "Command is not provided"
-
-    run_id = os.environ.get("RUNNABLE_TORCH_RUN_ID", "")
-    parameters_files = os.environ.get("RUNNABLE_TORCH_PARAMETERS_FILES", "")
-
-    process_run_id = (
-        run_id
-        + "-"
-        + os.environ.get("RANK", "")
-        + "-"
-        + "".join(random.choices(string.ascii_lowercase, k=3))
-    )
-    os.environ["TORCH_DISTRIBUTED_DEBUG"] = "DETAIL"
-
-    # In this subprocess there shoould not be any RUNNABLE environment variables
-    delete_env_vars_with_prefix("RUNNABLE_")
-
-    module_name, func_name = get_module_and_attr_names(command)
-    module = importlib.import_module(module_name)
-
-    callable_obj = getattr(module, func_name)
-
-    # The job runs with the default configuration
-    # ALl the execution logs are stored in .catalog
-    job = PythonJob(function=callable_obj)
-
-    config_content = {
-        "catalog": {"type": "file-system", "config": {"catalog_location": "proc_logs"}}
-    }
-
-    temp_config_file = Path("runnable-config.yaml")
-    with open(str(temp_config_file), "w", encoding="utf-8") as config_file:
-        yaml = YAML(typ="safe", pure=True)
-        yaml.dump(config_content, config_file)
-
-    job.execute(
-        parameters_file=parameters_files,
-        job_id=process_run_id,
-    )
-
-    # delete the temp config file
-    temp_config_file.unlink()
-
-    from runnable.context import run_context
-
-    job_log = run_context.run_log_store.get_run_log_by_id(run_id=run_context.run_id)
-
-    if job_log.status == defaults.FAIL:
-        raise Exception(f"Job {process_run_id} failed")