Merge pull request #194357 from BillHuang2001/init/ray

pkgs/development/python-modules/aiorwlock/default.nix

@@ -0,0 +1,31 @@
+{ buildPythonPackage
+, fetchPypi
+, lib
+, pytest-asyncio
+, pytest-cov
+, pytestCheckHook
+}:
+
+buildPythonPackage rec {
+  pname = "aiorwlock";
+  version = "1.3.0";
+
+  src = fetchPypi {
+    inherit pname version;
+    sha256 = "sha256-g/Eth99LlyiguP2hdWWFqw1lKxB7q1nGCE4bGtaSq0U=";
+  };
+
+  checkInputs = [
+    pytestCheckHook
+    pytest-asyncio
+  ];
+
+  pythonImportsCheck = [ "aiorwlock" ];
+
+  meta = with lib; {
+    description = "Read write lock for asyncio";
+    homepage = "https://github.com/aio-libs/aiorwlock";
+    license = licenses.asl20;
+    maintainers = with maintainers; [ billhuang ];
+  };
+}

pkgs/development/python-modules/gpustat/default.nix

@@ -0,0 +1,44 @@
+{ buildPythonPackage
+, blessed
+, fetchPypi
+, lib
+, mockito
+, nvidia-ml-py
+, psutil
+, pytest-runner
+, pythonRelaxDepsHook
+, pytestCheckHook
+}:
+
+buildPythonPackage rec {
+  pname = "gpustat";
+  version = "1.0.0";
+
+  src = fetchPypi {
+    inherit pname version;
+    sha256 = "sha256-WB6P+FjDLJWjIruPA/HZ3D0Xe07LM93L7Sw3PGf04/E=";
+  };
+
+  nativeBuildInputs = [ pythonRelaxDepsHook ];
+  pythonRelaxDeps = [ "nvidia-ml-py" ];
+
+  propagatedBuildInputs = [
+    blessed
+    nvidia-ml-py
+    psutil
+  ];
+
+  checkInputs = [
+    mockito
+    pytestCheckHook
+  ];
+
+  pythonImportsCheck = [ "gpustat" ];
+
+  meta = with lib; {
+    description = "A simple command-line utility for querying and monitoring GPU status";
+    homepage = "https://github.com/wookayin/gpustat";
+    license = licenses.mit;
+    maintainers = with maintainers; [ billhuang ];
+  };
+}

pkgs/development/python-modules/opencensus/default.nix

@@ -0,0 +1,55 @@
+{ buildPythonPackage
+, fetchPypi
+, lib
+, python
+, unittestCheckHook
+, google-api-core
+}:
+
+let
+  opencensus-context = buildPythonPackage rec {
+    pname = "opencensus-context";
+    version = "0.1.3";
+
+    checkInputs = [ unittestCheckHook ];
+
+    src = fetchPypi {
+      inherit pname version;
+      sha256 = "sha256-oDEIw8ENjIC7Xd9cih8DMWH6YZcqmRf5ubOhhRfwCIw=";
+    };
+  };
+in
+buildPythonPackage rec {
+  pname = "opencensus";
+  version = "0.11.0";
+
+  src = fetchPypi {
+    inherit pname version;
+    sha256 = "sha256-AmIWq6uJ2U2FBJLz3GWVAFXsT4QRX6bHvq/7pEo0bkI=";
+  };
+
+  buildInputs = [
+    # opencensus-context is embedded in opencensus
+    opencensus-context
+  ];
+
+  propagatedBuildInputs = [
+    google-api-core
+  ];
+
+  postInstall = ''
+    ln -sf ${opencensus-context}/${python.sitePackages}/opencensus/common/runtime_context \
+      $out/${python.sitePackages}/opencensus/common/
+  '';
+
+  checkInputs = [ unittestCheckHook ];
+
+  pythonImportsCheck = [ "opencensus" ];
+
+  meta = with lib; {
+    description = "A stats collection and distributed tracing framework";
+    homepage = "https://github.com/census-instrumentation/opencensus-python";
+    license = licenses.asl20;
+    maintainers = with maintainers; [ billhuang ];
+  };
+}

pkgs/development/python-modules/ray/binary-hashes.nix

@@ -0,0 +1,11 @@
+{
+  cp38 = {
+    sha256 = "da8adfa33ff54bc61cfe80334a0ee889e0060918db6ff9215aebe32e98b1f939";
+  };
+  cp39 = {
+    sha256 = "cab13346650f88171b3f348ed352f04695b96d1ab1090ed3b80bdc93e897dbd4";
+  };
+  cp310 = {
+    sha256 = "bcf3bff9517d77ea6c98592fa16e1cfb8bc0cfa345d3be69729bfa9c5bd78a7c";
+  };
+}

pkgs/development/python-modules/ray/default.nix

@@ -0,0 +1,166 @@
+{ aiohttp
+, aiohttp-cors
+, aiorwlock
+, aiosignal
+, attrs
+, autoPatchelfHook
+, buildBazelPackage
+, buildPythonPackage
+, fetchPypi
+, click
+, cloudpickle
+, colorama
+, colorful
+, cython
+, dm-tree
+, fastapi
+, filelock
+, frozenlist
+, fsspec
+, gpustat
+, grpc
+, grpcio
+, gym
+, jsonschema
+, lib
+, lz4
+, matplotlib
+, msgpack
+, numpy
+, opencensus
+, packaging
+, pandas
+, py-spy
+, prometheus-client
+, protobuf3_20
+, psutil
+, pyarrow
+, pydantic
+, python
+, pythonAtLeast
+, pythonOlder
+, pythonRelaxDepsHook
+, pyyaml
+, redis
+, requests
+, scikitimage
+, scipy
+, setproctitle
+, smart-open
+, starlette
+, stdenv
+, tabulate
+, tensorboardx
+, uvicorn
+, virtualenv
+}:
+
+let
+  pname = "ray";
+  version = "2.0.0";
+in
+buildPythonPackage rec {
+  inherit pname version;
+  format = "wheel";
+
+  disabled = pythonOlder "3.8" || pythonAtLeast "3.11";
+
+  src =
+    let
+      pyShortVersion = "cp${builtins.replaceStrings ["."] [""] python.pythonVersion}";
+      binary-hash = (import ./binary-hashes.nix)."${pyShortVersion}";
+    in
+    fetchPypi ({
+      inherit pname version format;
+      dist = pyShortVersion;
+      python = pyShortVersion;
+      abi = pyShortVersion;
+      platform = "manylinux2014_x86_64";
+    } // binary-hash);
+
+  passthru.optional-dependencies = rec {
+    data-deps = [
+      pandas
+      pyarrow
+      fsspec
+    ];
+
+    serve-deps = [
+      aiorwlock
+      fastapi
+      pandas
+      starlette
+      uvicorn
+    ];
+
+    tune-deps = [
+      tabulate
+      tensorboardx
+    ];
+
+    rllib-deps = tune-deps ++ [
+      dm-tree
+      gym
+      lz4
+      matplotlib
+      scikitimage
+      pyyaml
+      scipy
+    ];
+
+    air-deps = data-deps ++ serve-deps ++ tune-deps ++ rllib-deps;
+  };
+
+  nativeBuildInputs = [
+    autoPatchelfHook
+    pythonRelaxDepsHook
+  ];
+
+  pythonRelaxDeps = [ "grpcio" "click" "protobuf" ];
+
+  propagatedBuildInputs = [
+    attrs
+    aiohttp
+    aiohttp-cors
+    aiosignal
+    click
+    cloudpickle
+    colorama
+    colorful
+    cython
+    filelock
+    frozenlist
+    gpustat
+    grpcio
+    jsonschema
+    msgpack
+    numpy
+    opencensus
+    packaging
+    py-spy
+    prometheus-client
+    protobuf3_20
+    psutil
+    pydantic
+    pyyaml
+    requests
+    setproctitle
+    smart-open
+    virtualenv
+  ];
+
+  postInstall = ''
+    chmod +x $out/${python.sitePackages}/ray/core/src/ray/{gcs/gcs_server,raylet/raylet}
+    ln -sf ${redis}/bin/redis-server $out/${python.sitePackages}/ray/core/src/ray/thirdparty/redis/src/redis-server
+  '';
+
+  pythonImportsCheck = [ "ray" ];
+
+  meta = with lib; {
+    description = "A unified framework for scaling AI and Python applications";
+    homepage = "https://github.com/ray-project/ray";
+    license = licenses.asl20;
+    maintainers = with maintainers; [ billhuang ];
+    platforms = [ "x86_64-linux" ];
+  };
+}

pkgs/top-level/python-packages.nix

@@ -433,6 +433,8 @@ in {
 
   aiorun = callPackage ../development/python-modules/aiorun { };
 
+  aiorwlock = callPackage ../development/python-modules/aiorwlock { };
+
   aiosenseme = callPackage ../development/python-modules/aiosenseme { };
 
   aiosenz = callPackage ../development/python-modules/aiosenz { };
@@ -4035,6 +4037,8 @@ in {
 
   gpsoauth = callPackage ../development/python-modules/gpsoauth { };
 
+  gpustat = callPackage ../development/python-modules/gpustat { };
+
   gpxpy = callPackage ../development/python-modules/gpxpy { };
 
   gpy = callPackage ../development/python-modules/gpy { };
@@ -6520,6 +6524,8 @@ in {
       openbabel = callPackage ../development/libraries/openbabel { inherit (self) python; };
   };
 
+  opencensus = callPackage ../development/python-modules/opencensus { };
+
   opencv3 = toPythonModule (pkgs.opencv3.override {
     enablePython = true;
     pythonPackages = self;
@@ -9622,6 +9628,8 @@ in {
 
   rawkit = callPackage ../development/python-modules/rawkit { };
 
+  ray = callPackage ../development/python-modules/ray { redis = pkgs.redis; };
+
   rbtools = callPackage ../development/python-modules/rbtools { };
 
   rcssmin = callPackage ../development/python-modules/rcssmin { };