hunk.code_window:     covered_o = set()
       config["log_level"] = "ERROR"
       first_error = None
       for a_name, action_space in ACTION_SPACES_TO_TEST.items():
           for o_name, obs_space in OBSERVATION_SPACES_TO_TEST.items():
  
hunk.labels: keep keep add keep keep
hunk.after_edit:     torch = config.get("use_pytorch", False)
  
hunk.file_path: rllib/tests/test_supported_spaces.py
hunk.type: add
file: package org.ray.runtime.task;
  
  import com.google.common.base.Preconditions;
  import java.util.ArrayList;
  import java.util.List;
  import java.util.concurrent.ConcurrentHashMap;
  import org.ray.api.exception.RayTaskException;
  import org.ray.api.id.ActorId;
  import org.ray.api.id.JobId;
  import org.ray.api.id.TaskId;
  import org.ray.api.id.UniqueId;
  import org.ray.runtime.AbstractRayRuntime;
  import org.ray.runtime.config.RayConfig;
  import org.ray.runtime.config.RunMode;
  import org.ray.runtime.functionmanager.JavaFunctionDescriptor;
  import org.ray.runtime.functionmanager.RayFunction;
  import org.ray.runtime.generated.Common.TaskType;
  import org.ray.runtime.object.NativeRayObject;
  import org.ray.runtime.object.ObjectSerializer;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * The task executor, which executes tasks assigned by raylet continuously.
   */
  public abstract class TaskExecutor {
  
    private static final Logger LOGGER = LoggerFactory.getLogger(TaskExecutor.class);
  
    // A helper map to help we get the corresponding executor for the given worker in JNI.
    private static ConcurrentHashMap<UniqueId, TaskExecutor> taskExecutors
        = new ConcurrentHashMap<>();
  
    protected final AbstractRayRuntime runtime;
  
    /**
     * The current actor object, if this worker is an actor, otherwise null.
     */
    protected Object currentActor = null;
  
    /**
     * The exception that failed the actor creation task, if any.
     */
    private Exception actorCreationException = null;
  
    protected TaskExecutor(AbstractRayRuntime runtime) {
      this.runtime = runtime;
      if (RayConfig.getInstance().runMode == RunMode.CLUSTER) {
        taskExecutors.put(runtime.getWorkerContext().getCurrentWorkerId(), this);
      }
    }
  
    public static TaskExecutor get(byte[] workerId) {
      return taskExecutors.get(new UniqueId(workerId));
    }
  
    protected List<NativeRayObject> execute(List<String> rayFunctionInfo,
        List<NativeRayObject> argsBytes) {
      JobId jobId = runtime.getWorkerContext().getCurrentJobId();
      TaskType taskType = runtime.getWorkerContext().getCurrentTaskType();
      TaskId taskId = runtime.getWorkerContext().getCurrentTaskId();
      LOGGER.debug("Executing task {}", taskId);
  
      List<NativeRayObject> returnObjects = new ArrayList<>();
      ClassLoader oldLoader = Thread.currentThread().getContextClassLoader();
      JavaFunctionDescriptor functionDescriptor = parseFunctionDescriptor(rayFunctionInfo);
      RayFunction rayFunction = null;
      try {
        // Find the executable object.
        rayFunction = runtime.getFunctionManager().getFunction(jobId, functionDescriptor);
        Thread.currentThread().setContextClassLoader(rayFunction.classLoader);
        runtime.getWorkerContext().setCurrentClassLoader(rayFunction.classLoader);
  
        // Get local actor object and arguments.
        Object actor = null;
        if (taskType == TaskType.ACTOR_TASK) {
          if (actorCreationException != null) {
            throw actorCreationException;
          }
          actor = currentActor;
  
        }
        Object[] args = ArgumentsBuilder.unwrap(argsBytes, rayFunction.classLoader);
        // Execute the task.
        Object result;
        if (!rayFunction.isConstructor()) {
          result = rayFunction.getMethod().invoke(actor, args);
        } else {
          result = rayFunction.getConstructor().newInstance(args);
        }
        // Set result
        if (taskType != TaskType.ACTOR_CREATION_TASK) {
          if (taskType == TaskType.ACTOR_TASK) {
            // TODO (kfstorm): handle checkpoint in core worker.
            maybeSaveCheckpoint(actor, runtime.getWorkerContext().getCurrentActorId());
          }
          if (rayFunction.hasReturn()) {
            returnObjects.add(ObjectSerializer.serialize(result));
          }
        } else {
          // TODO (kfstorm): handle checkpoint in core worker.
          maybeLoadCheckpoint(result, runtime.getWorkerContext().getCurrentActorId());
          currentActor = result;
        }
        LOGGER.debug("Finished executing task {}", taskId);
      } catch (Exception e) {
        LOGGER.error("Error executing task " + taskId, e);
        if (taskType != TaskType.ACTOR_CREATION_TASK) {
          boolean hasReturn = rayFunction != null && rayFunction.hasReturn();
          boolean isCrossLanguage = functionDescriptor.signature.equals("");
          if (hasReturn || isCrossLanguage) {
            returnObjects.add(ObjectSerializer
                .serialize(new RayTaskException("Error executing task " + taskId, e)));
          }
        } else {
          actorCreationException = e;
        }
      } finally {
        Thread.currentThread().setContextClassLoader(oldLoader);
        runtime.getWorkerContext().setCurrentClassLoader(null);
      }
      return returnObjects;
    }
  
    private JavaFunctionDescriptor parseFunctionDescriptor(List<String> rayFunctionInfo) {
      Preconditions.checkState(rayFunctionInfo != null && rayFunctionInfo.size() == 3);
      return new JavaFunctionDescriptor(rayFunctionInfo.get(0), rayFunctionInfo.get(1),
          rayFunctionInfo.get(2));
    }
  
    protected abstract void maybeSaveCheckpoint(Object actor, ActorId actorId);
  
    protected abstract void maybeLoadCheckpoint(Object actor, ActorId actorId);
  }
  