Repository: ray-project/ray
File: release/ray_release/command_runner/anyscale_job_runner.py (Python)

Code before:

    def _copy_script_to_working_dir(self, script_name):
        script = os.path.join(os.path.dirname(__file__), f"_{script_name}")
        shutil.copy(script, script_name)

    def prepare_remote_env(self):
        self._copy_script_to_working_dir("anyscale_job_wrapper.py")
        self._copy_script_to_working_dir("wait_cluster.py")
        self._copy_script_to_working_dir("prometheus_metrics.py")

    def run_prepare_command(
        self, command: str, env: Optional[Dict] = None, timeout: float = 3600.0
    ):
        self.prepare_commands.append((command, env, timeout))

    def wait_for_nodes(self, num_nodes: int, timeout: float = 900):
        self._wait_for_nodes_timeout = timeout
        self.job_manager.cluster_startup_timeout += timeout

        # Give 30 seconds more to account for communication
        self.run_prepare_command(
            f"python wait_cluster.py {num_nodes} {timeout}", timeout=timeout + 30
        )

    def _handle_command_output(self, job_state: int, raise_on_timeout: bool = True):
        if job_state == -1:
            raise JobOutOfRetriesError(
                "Job returned non-success state: 'FAILED' "
                "(command has not been ran or no logs could have been obtained)."
            )

        if job_state == -2:
            raise JobBrokenError("Job state is 'UNKNOWN'.")

        # First try to obtain the output.json from S3.
        # If that fails, try logs.
        try:
            output_json = self.fetch_output()
        except Exception:
            logger.exception("Exception when obtaining output from S3.")
            try:
                logs = self.get_last_logs()
                output_json = re.search(r"### JSON \|([^\|]*)\| ###", logs)
                output_json = json.loads(output_json.group(1))
            except Exception:
                output_json = None

        workload_status_code = None
        if output_json:
            logger.info(f"Output: {output_json}")
            workload_status_code = output_json["return_code"]

Review comment (bug):
**non-blocking**: Instead of `-1` and `0`, I wonder if there'd be a way to immediately surface where these are from. E.g. 

```
if job_state == job_status_to_return_code[JobState.FAILED]
```

Mainly food for thought.

Code after:
        # artifact related
        # user provided path to where they write the artifact to.
        self._artifact_path = artifact_path
        self._artifact_uploaded = artifact_path is not None

    def _copy_script_to_working_dir(self, script_name):
        script = os.path.join(os.path.dirname(__file__), f"_{script_name}")
        shutil.copy(script, script_name)

    def prepare_remote_env(self):
        self._copy_script_to_working_dir("anyscale_job_wrapper.py")
        self._copy_script_to_working_dir("wait_cluster.py")
        self._copy_script_to_working_dir("prometheus_metrics.py")

    def run_prepare_command(
        self, command: str, env: Optional[Dict] = None, timeout: float = 3600.0
    ):
        self.prepare_commands.append((command, env, timeout))

    def wait_for_nodes(self, num_nodes: int, timeout: float = 900):
        self._wait_for_nodes_timeout = timeout
        self.job_manager.cluster_startup_timeout += timeout

        # Give 30 seconds more to account for communication
        self.run_prepare_command(
            f"python wait_cluster.py {num_nodes} {timeout}", timeout=timeout + 30
        )

    def _handle_command_output(
        self, job_return_code: int, raise_on_timeout: bool = True
    ):
        if job_return_code == JOB_FAILED:
            raise JobOutOfRetriesError(
                "Job returned non-success state: 'FAILED' "
                "(command has not been ran or no logs could have been obtained)."
            )

        if job_return_code == JOB_STATE_UNKNOWN:
            raise JobBrokenError("Job state is 'UNKNOWN'.")

        # First try to obtain the output.json from S3.
        # If that fails, try logs.
        try:
            output_json = self.fetch_output()
        except Exception:
            logger.exception("Exception when obtaining output from S3.")
            try:
                logs = self.get_last_logs()
                output_json = re.search(r"### JSON \|([^\|]*)\| ###", logs)
                output_json = json.loads(output_json.group(1))
            except Exception: